If I understand you correctly, you would like to select from the last dimension using an index tensor with values in [0, 2].
The same operation should be performed for all examples?
If so, you could try the following code:
a = torch.randn(20, 10, 3)

idx = torch.zeros(10).random_(3).long()
idx = idx.unsqueeze(0).repeat(20, 1)
 
b = a.gather(2, idx.view(20, 10, 1))
print(b.shape)

I’ve just used 20 for the batch dimension.
Let me know, if this works for you or if I misunderstood your question.