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.