Accessing tensor elements using another "index" tensor

1

Let’s say I have a tensor A of shape (BatchxDim1xDim2x…) and a tensor B of shape (Batch)
B = Tensor([1,3,7,2,…])
that encodes indices that I am interested in.
What is the best way to obtain a tensor C such that
C[i,j,k,…] = A[i,B[i],j,k,…] ?

2

Hi,

The torch.gather function is built just for that.
You might have to expand your B to get the behavior you want here.

1 Like
3

Yes, this works. Thanks a lot!