How to extract the Probability of specific class from the softmax output

Hi all,

I have a question that how to extract the probability of specific class from the softmax output?
Given specific class labels, how can we extract the probability ?

Thanks.

For example, we use CNN to classify two classes, and its outputs are as follows.
0.4 0.6
0.5 0.5
0.2 0.8
specific class labels:
1
0
1
So, we get the result:
0.6
0.5
0.8

You could use torch.gather for this:

output = Variable(torch.randn(10, 3))
prob = F.softmax(output, dim=1)
classes = Variable(torch.LongTensor(10, 1).random_(0, 3))
class_prob = torch.gather(prob, 1, classes)

It works well.
Thank you very much.