Super basic question, but I am having problems with it.
Essentially I have a tensor called probs with shape (batch_size, classes) let’s say (4, 3), and another tensor called labs with shape (batch_size), let’s say (4). Let’s assume that the values of the tensors are:
a = tensor([[0.1, 0.2, 0.3, 0.2, 0.2],
            [0.4, 0.3, 0.2, 0.1, 0   ],
            [0.1, 0.4, 0.3, 0.1, 0.1],
            [0.2, 0.2, 0.2, 0.1, 0.3]]

and
labs = tensor([2, 1, 0, 3])

I want the result to be:
result = tensor([0.3,
                [0.3,
                [0.1, 
                [0.1])

essentially getting the 2nd (zero indexing) value from the zeroth row, first value from the second row, zeroth value from the second row and the third value from the third row.
Does anyone know what is the best way of achieving this (not using for loops)? I thought that
probs[labs.unsqueeze(1)]

should do the trick, but apparently not. Also, tried with gather, but no luck so far.