Select Rows of 2D Tensor

I have a tensor of indices:

tensor([835, 299, 399, 975])

and an embedding matrix

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.

I’ve found it, it can be done via

A.index_select(0, user_idx)