nn.NLLLoss expects a model output in the shape [batch_size, nb_classes] containing log probabiltiies and a target in the shape [batch_size] as a LongTensor containing class indices in the range [0, nb_classes-1].
Based on the error message, I guess your target shape is wrong. If you are using a one-hot encoded target, use target = target.argmax(dim=1) to create the class indices.
1 Like