PyTorch训练时出现RuntimeError：张量无梯度及grad_fn问题求助

核心问题 #

报错根源是**argmax属于不可微分操作**，执行该操作后生成的张量会彻底丢失梯度传递信息，导致后续计算的loss没有grad_fn，无法触发反向传播。

同时代码存在两处次要错误：

• 测试阶段计算test_loss时误用了训练集标签y_train，应替换为测试集标签y_test
• 分类任务中，直接对argmax输出的离散类别索引计算损失不符合逻辑（损失函数通常接收模型输出的logits或概率分布，而非离散索引）

修复步骤 #

1. 修正训练与测试的核心逻辑 #

移除损失计算链路中的argmax操作，仅在计算准确率时使用该操作生成预测结果：

"""Training"""
Epochs = 100

for epoch in range(Epochs):
    model.train()

    train_logits = model(X_train)
    # 直接用模型输出的logits计算损失（CrossEntropyLoss内部会自动处理softmax）
    loss = loss_fn(train_logits, y_train)
    # 仅计算准确率时生成预测结果
    train_preds = torch.softmax(train_logits, dim=1).argmax(dim=1)
    train_accu = accuracy(y_train, train_preds)

    optimiser.zero_grad()
    loss.backward()
    optimiser.step()

    # 测试阶段
    model.eval()
    with torch.inference_mode():
        test_logits = model(X_test)
        test_preds = torch.softmax(test_logits, dim=1).argmax(dim=1)
        # 修正测试集标签引用
        test_loss = loss_fn(test_logits, y_test)
        test_acc = accuracy(y_test, test_preds)
    
    if epoch%10 == 0:
        print(f'Epoch:{epoch} | Train loss: {loss:.4f} | Training acc:{train_accu:.4f} | Test Loss: {test_loss:.4f} | Test accu: {test_acc:.4f}')

2. 额外说明 #

• 若使用NLLLoss作为损失函数，需先对logits执行log_softmax，但同样不能在损失计算前使用argmax
• Sequential模型默认会为所有可训练参数开启requires_grad=True，无需额外设置，除非手动冻结了部分参数

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