Solved by ptrblck in post #3 
              
              
                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)