Torch select index for mult dimensional vector
Suppose I have the following data:
import torch
torch.manual_seed(42)
logits = torch.randn(2, 5, 3)
idx = torch.randint(0, 3, (2, 5))

What I want to do is the following:
[[logits[i,j,idx[i][j]] for j in range(len(idx[i]))] for i in range(len(idx))]

However, this is clearly inefficient.
The closest that I came to this was doing this, but this again looks pretty ugly:
new_idx = torch.stack([idx]*logits.shape[-1], dim=-1)
logits.gather(dim=-1, index=new_idx)

where in the above case the desired output is replicated 3 times.
In terms of practical use case, I'm thinking of a language model where the logits are of shape (batch_size, sequence_len, vocabulary), and the indices are simply (batch_size, sequence_len). Apologies if this has been asked before but I couldn't find anything.

same with this answer
torch.gather(logits, -1, idx.unsqueeze(-1))
&gt;&gt;&gt; [[logits[i,j,idx[i][j]] for j in range(len(idx[i]))] for i in range(len(idx))]
[[tensor(0.9007), tensor(0.6784), tensor(-0.0431), tensor(-1.4036), tensor(-0.7279)], [tensor(-0.2168), tensor(1.7174), tensor(-0.4245), tensor(0.9956), tensor(-1.2742)]]

&gt;&gt;&gt; torch.gather(logits, -1, idx.unsqueeze(-1))
tensor([[[ 0.9007],
         [ 0.6784],
         [-0.0431],
         [-1.4036],
         [-0.7279]],

        [[-0.2168],
         [ 1.7174],
         [-0.4245],
         [ 0.9956],
         [-1.2742]]])

