To get the prediction, you can use torch.argmax(output, 1).
The logits will give you the same prediction as the softmax output.
If you would like to see the “probabilities” of your output, you could also use
probs = F.softmax(output, 1)

but don’t pass probs to your criterion, as this might result in bad training.