Efficient way to obtain values from specified indexes

I have a tensor A of shape(L, W). I also have another tensor B of shape(L, 1). I want to use B as indexes to select from A, which would result in a tensor C of shape(L). For example, say A is:[ [0, 1], [2, 3], [5, 6] ], and B is[ [1], [1], [0] ], I want to make C contains: [1, 3, 5]. I didn’t find any function that suits my needs. Currently, I’m achieving it by doing this:

        mask: torch.Tensor = torch.zeros_like(A, dtype=torch.uint8)
        mask.scatter_(1, B, 1)
        C = A.masked_select(mask)

I feel that this method isn’t very efficient as it first needs to convert B into one-hot and then selecting. So is there a better way?

Try this

import torch
A = torch.tensor([[0, 1], [2, 3], [5, 6]]).float()
B = torch.tensor([[1], [1], [0]]).long()
print(torch.gather(A, 1, B))

It works, thx. I was long troubled by what does torch.gather do :rofl: