I assume you are working on a multi-class classification use case with nn.CrossEntropyLoss as the criterion.
If that’s the case, you would have to make sure that the model output has the shape [batch_size, nb_classes], while the target should have the shape [batch_size] containing the class indices in the range [0, nb_classes-1]. For 5 classes, your target should contain the values [0, 1, 2, 3, 4], while it seems you are passing a 5 to the criterion.