PyTorch运行前馈NN报张量设备不匹配RuntimeError排查

• 张量设备不匹配：代码中仅对输入images做了.to(device)迁移，和预测结果一同传入损失函数的labels张量始终留在CPU上。交叉熵损失计算要求预测张量、目标标签张量必须位于同一设备，这是触发设备不匹配报错的直接原因。
  你提到当前环境执行torch.cuda.is_available()返回False却出现cuda:0与CPU不匹配的提示，是因为之前运行代码时CUDA处于可用状态，Python内核缓存了驻留在GPU上的模型/张量残留，重启内核后设备选择逻辑回退到CPU，但核心的标签未迁移设备的bug始终存在。
  补充说明：损失函数nn.CrossEntropyLoss、Adam优化器本身没有可迁移的参数，调用.to(device)没有实际作用，无法解决张量不匹配问题。
• 训练循环逻辑错误：现有代码完全缺失优化器参数更新的optimizer.step()调用，还把optimizer.zero_grad()放在了loss.backward()之后，会直接把反向传播计算得到的梯度全部清零，就算解决设备报错，模型参数也无法正常更新。
• 测试集加载逻辑错误：初始化test_loader时错误传入了训练集train_dataset，后续测试环节会直接复用训练数据，无法验证模型泛化能力。另外代码中给torchvision.transforms起的别名为transformers，容易和HuggingFace的同名第三方库混淆，属于不规范写法。

按照以下规则修改代码即可：

• 取到每个批次的images和labels后，将两个张量都迁移到指定device
• 调整训练循环的执行顺序：每轮迭代先执行梯度清零，再做前向传播算损失、反向传播计算梯度，最后调用优化器更新参数
• 修正test_loader的入参，传入测试集test_dataset，修正transforms的别名避免混淆

修复后的完整可运行代码如下：

import torch
import torch.nn as nn
import torchvision
import torchvision.transforms as transforms

# 超参数
input_size = 784
hidden_size = 100
num_classes = 10
num_epochs = 2
learning_rate = 0.001
batch_size = 100

# 计算设备选择
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

# 加载数据
train_dataset = torchvision.datasets.MNIST('.', train=True, transform=transforms.ToTensor(), download=True)
test_dataset = torchvision.datasets.MNIST('.', train=False, transform=transforms.ToTensor(), download=False)

train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
# 修复：测试集加载器传入test_dataset
test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=batch_size, shuffle=False)


class FeedforwardNN(nn.Module):
    def __init__(self, input_size, hidden_size, num_classes):
        super(FeedforwardNN, self).__init__()
        self.l1 = nn.Linear(input_size, hidden_size)
        self.r1 = nn.ReLU()
        self.l2 = nn.Linear(hidden_size, num_classes)

    def forward(self, x):
        out = self.l1(x)
        out = self.r1(out)
        out = self.l2(out)
        return out


model = FeedforwardNN(input_size, hidden_size, num_classes).to(device)

# 损失函数与优化器
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)


# 训练循环
n_total_steps = len(train_loader)

for epoch in range(num_epochs):
    for i, (images, labels) in enumerate(train_loader):
        # 修复：images和labels都迁移到指定设备
        images = images.reshape(-1, 28*28).to(device)
        labels = labels.to(device)
        
        # 修复：调整训练步骤顺序
        # 1. 梯度清零
        optimizer.zero_grad()
        # 2. 前向传播
        y_pred = model(images)
        # 3. 计算损失
        loss = criterion(y_pred, labels)
        # 4. 反向传播
        loss.backward()
        # 5. 参数更新
        optimizer.step()

        # 可选：打印训练日志
        if (i+1) % 100 == 0:
            print(f'Epoch [{epoch+1}/{num_epochs}], Step [{i+1}/{n_total_steps}], Loss: {loss.item():.4f}')

内容的提问来源于stack exchange，提问作者Salih