torch.nn.functional.softmax(output, dim = 1) will yield an output, where the probabilities for each pixel sum to 1 in the class dimension (dim1).
This probability tensor can be used as a sanity check or for visualization purposes.
However, note that e.g. nn.CrossEntropyLoss expects raw logits as the model’s output, since internally F.log_softmax(output, 1) and nn.NLLLoss will be used, so that you shouldn’t use the softmax method manually.