tcsn_wty:

But if I understand correctly, you meant that I need to do the permutation before log_softmax called, am I correct?


If you want to keep F.log_softmax(x, dim=1), then yes.
Otherwise use dim=2, if you want to permute the tensor afterwards.



 tcsn_wty:

As for loss function, if I’m using nn.CrossEntropyLoss, do you meant that there’s no need to apply F.log_softmax in forward() function?


Yes, since it will be applied twice at the moment (in your forward and inside nn.CrossEntropyLoss).