I currently have a tensor X of shape (2,3,4,10), and an indexing vector Y of shape (4):

import torch
X=tensor([[[[8, 2, 8, 5],
          [3, 7, 4, 0]],

         [[4, 5, 7, 4],
          [8, 3, 9, 5]]],


        [[[5, 2, 9, 3],
          [6, 4, 5, 4]],

         [[7, 3, 3, 7],
          [6, 3, 8, 9]]]])
Y=tensor([1,1,0,1])

I would like to index X with Y, such that I arrive at a tensor of shape (3,4,10), which, if manually done, should look as follows (with the first dimension corresponding to the indices in Y):

torch.stack([X[1][0][0], X[1][0][1], X[0][1][0], X[1][1][1]])

Which gives:

tensor([[5, 2, 9, 3],
        [6, 4, 5, 4],
        [4, 5, 7, 4],
        [6, 3, 8, 9]])

I currently use a for loop to do this, but this is not ideal, and was wondering how to get to the desired result tensor in PyTorch from X and Y. Thanks a lot!

I tried a for loop to index into it, which works, but this is very slow.