Solved by ptrblck in post #2 
              
              
                If your model returns logits, you could call F.softmax on them, which will yield the class probabilities. 
Example: 
import torch.nn.functional as F

model = nn.Linear(10, 5)
x = torch.randn(1, 10)
output = model(x)
prob = F.softmax(output, dim=1)