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?