Oli:

a softmax whilst your loss function, cross_entropy, doesn’t want this. See the docs about cross entropy and NLLLoss


 I totally missed the loss function using F.cross_entropy at the bottom of the code snippet. Sorry about that! Disregard pedantic note about label shapes above.
I do think the function is a bit confusing in this case. The docs say:

input has to be a Tensor of size either (minibatch,C)(minibatch, C)(minibatch,C) or (minibatch,C,d1,d2,…,dK)(minibatch, C, d_1, d_2, …, d_K)(minibatch,C,d1​,d2​,…,dK​) with K≥1K \geq 1K≥1 for the K-dimensional case (described later).

and

This criterion expects a class index in the range [0,C−1][0, C-1][0,C−1] as the target for each value of a 1D tensor of size minibatch;

So the network is still supposed to be returning probabilities over a dimension with rank = (your number of classes). As opposed to what you described:



 Rojin:

the output is a 1D tensor with the index of the label ex: tensor([2,5,31,…,7])


But the label is as you describe, a tensor of integer indices, where each entry is the integer class label that corresponds to the ground truth for the matching input.
Since the shapes all seem to match (32 on the output), it sounds as @Oli mentioned an activation function problem.