为什么PyTorch训练分类模型时第一个epoch后损失始终返回0？

• 标签维度不匹配：你当前的labels张量是二维结构（形状为[batch_size, 1]，示例中为[64, 1]），而nn.CrossEntropyLoss()要求输入的目标张量必须是1维（形状为[batch_size]），这就是你直接传labels会抛出RuntimeError: 0D or 1D target tensor expected, multi-target not supported错误的直接原因。
• 损失始终为0的原因：你当前使用的torch.max(labels, 1)[1]代码逻辑错误。因为labels在dim=1维度上只有1个元素，调用torch.max在dim=1取索引时，返回的所有位置的索引值都是0，相当于你把所有样本的标签强制替换成了类别0。模型只要快速学会把所有样本都预测为类别0，交叉熵损失就会降到接近0，这就是损失异常为0的根源。

只需要在计算损失前，把二维的标签张量压缩为1维即可，修改损失计算的代码行：

# 去掉dim=1维度上长度为1的冗余维度，得到形状为[batch_size]的1维标签
loss = loss_fn(outputs, labels.squeeze(1))

删除你原有错误的torch.max(labels, 1)[1]写法即可正常训练。

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