You cant use nn.CrossEntropyLoss with softmax in pytorch.
use as per details given by @ptrblck .
If you really need softmax probabilities at the end then you can get using torch.exp(output)
and if you really want to train with softmax function than use custom cross entropy loss.
something like
def custom_categorical_cross_entropy(y_pred, y_true):
y_pred = torch.clamp(y_pred, 1e-9, 1 - 1e-9)
return -(y_true * torch.log(y_pred)).sum(dim=1).mean()