Solved by ptrblck in post #2 
              
              
                torch.gather should do the job: 
outputs = torch.FloatTensor([[0.3, 0.7], [0, 1], [0.4, 0.6]])
labels = torch.Tensor([0, 1, 1]).long()
print(outputs.gather(1, labels.unsqueeze(1)))
>tensor([[0.3000],
        [1.0000],
        [0.6000]])