If you are using torch.nn.CrossEntropyLoss() then you don’t need a softmax output layer on your model.
So it would just be

output = model(input) #logit output
criterion =torch.nn.CrossEntropyLoss()
loss = criterion(output, target)