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?