nn.CrossEntropyLoss expects raw logits in the shape [batch_size, nb_classes, *] so you should not apply a softmax activation on the model output. The class dimension should be in dim1 in the model output.
1 Like