Hi, thanks for your reply!
I tried your solution, but for more than one batches, ie like
x = torch.tensor( [[ [0.1, 0.2], [0.3, 0.4], [0.5, 0.6]], [ [0.7, 0.8], [0.9, 1.0], [1.1, 1.2]]]) # shape = (2,3,2)
i = torch.tensor([[1,0,1],[1,1,1]]) # shape = (2,3)

x[:, torch.arange(x.size(1)), i] broadcasts the last dim and gives
tensor([[[0.2000, 0.3000, 0.6000],
         [0.2000, 0.4000, 0.6000]],

        [[0.8000, 0.9000, 1.2000],
         [0.8000, 1.0000, 1.2000]]])
# shape (2,2,3)

where it should be
tensor([[0.2000, 0.3000, 0.6000],
         [0.8000, 1.0000, 1.2000]])
#shape (2,3)

was trying to index it again but i guess i still need a loop?