Hi,
This can be achieved by simply using nn.Softmax(). You can add it in your model definition to get index and probs at the same time. Right after this line:



 rizal_alfarizi:

Linear(128, num_class)
Softmax()


Or let’s say you have outputs (logits) and you are looking for probs here in your code:



 rizal_alfarizi:

output = model(input)


You can achieve it by same trick:
softmax = nn.Softmax()
output = model(input)
probs = softmax(output)

Bests