在使用pytorch交叉熵损失函数loss = nn.CrossEntropyLos(v, targets)时发现报错

IndexError: Dimension out of range (expected to be in range of [-1, 0], but got 1)。

排查后发现模型预测值V在训练阶段最后一个step时，tensor维度从二维变成一维，导致出错。

例：训练集个数81，bs设置为4时，最后一个step 只剩1张图片。最后step时tensor维度从二维变成一维。

于是添加以下判断语句解决问题。如果有其他方法，可以交流一下。

1.

 v = net(inputs)
 if len(v.shape)==1:  #方式出现训练最后一个step时，出现v是一维的情况
     v=torch.unsqueeze(v,0)  
 loss = loss_func(v, targets)

2.可以在DataLoader中drop_last参数设置False