In your current forward implementation you are passing x to F.softmax so that the actual model will not be used.
Also, nn.CrossEntropyLoss expects raw logits, as internally F.log_softmax will be applied, so remove the F.softmax and return out directly.
If you are still facing a shape mismatch error, could you print the shape of the data tensor before passing it to the model?