Select Rows of 2D Tensor

1

I have a tensor of indices:

user_idx
torch.Size([4])
tensor([835, 299, 399, 975])

and an embedding matrix

A
torch.Size([1000, 10])

Now I’d like to select the rows 835, 299, 399, 975 from matrix A. How can I do that such that autograd realizes that I’m doing this? I’ve looked at the gather function but I can’t understand how it works.

2

I’ve found it, it can be done via

A.index_select(0, user_idx)