PyTorch计算两个张量批次交叉熵时出现维度不匹配报错该如何解决？

首先明确两个张量的形状：

• 输入预测张量x的形状为[1, 3, 5]，三个维度分别对应batch_size=1、序列长度/样本数=3、类别数=5
• 标签张量y的形状为[1, 3]，对应每个batch下3个样本的分类标签

PyTorch的nn.CrossEntropyLoss默认将输入张量的第2个维度（索引为1的维度）作为类别维度。你当前的类别维度在最后一维（索引为2），所以损失函数会误将3识别为类别数，认为最后一维长度5是需要匹配的标签维度，因此抛出期望target形状为[1,5]的错误。

有两种常用修正方式，效果完全等价：

1. 调整输入张量的维度顺序，把类别维度移到第2位
2. 显式指定类别维度的位置（PyTorch 1.10及以上版本支持）

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

import torch
import torch.nn as nn

x = torch.tensor([[[ 2.1137, -1.3133,  0.7930,  0.3330,  0.9407],
        [-0.8380, -2.0299, -1.1218,  0.3150,  0.4797],
        [-0.7439,  0.0753, -0.1121,  0.0096, -1.2621]]])

y = torch.tensor([[1,2,3]])

# 方案1：调整维度顺序，兼容性更好
loss = nn.CrossEntropyLoss()(x.permute(0,2,1), y)
print(loss)

# 方案2：显式指定类别维度位置，代码更简洁
loss = nn.CrossEntropyLoss()(x, y, dim=2)
print(loss)

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