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?