The shapes look alright, if you add the discussed permute.
However, the F.log_softmax operation should be applied in the class dimension, so in your case dim2 (or permute the tensor before the operation).
Also, since you are using F.log_softmax, you should use nn.NLLLoss as the criterion.
nn.CrossEntropyLoss will apply F.log_softmax and nn.NLLLoss internally.