Hi, apologies if this has been asked before.

I have a tensor `x`

of shape `(batch_size, seq_len, vocab_size)`

, say shape `(1, 3, 2)`

. I have a tensor `i`

which is `(batch_size, vocab_indices)`

, say shape `(1,3)`

which are the indices of the vocab size dimension in `x`

. I want to get values from the last dimension (vocab_size).

For example,

```
x = torch.tensor( [[ [0.1, 0.2], [0.3, 0.4], [0.5, 0.6] ]]) # shape = (1,3,2)
i = torch.tensor([[1,0,1]]) # shape = (1,3)
```

I want to get

```
y = torch.tensor([[0.2,0.3,0.6]]) #shape = (1,3)
```

Is there a clean way to do this other than using a for loop? Thank you very much!