Hi,
I have a 3d tensor sized (batch_size, S, 5).
I also have a list containing a list for each batch, that contains indices.
For example, if the batch size is 2, the list looks like this:
[[idx0_0, idx0_1, …], [idx1_0, idx1_1, …]] (where each idx is between 0 to S).
I want to get a tensor that contains, for each batch, the 3rd dimension’s stuff from the original tensor, only from the indices that appear in the list of the same batch, while keeping the gradient of the original tensor.
For example, if the original tensor is
[[[A],
  [B], 
  [C]]]

(batch_size = 1, S = 3, A, B, C are sized 5)
and the list is [[0, 2]]
The result will be:
[[[A], 
  [C]]]

and that it’ll be possible to backpropagate through this tensor.
I heard about the gather function but i’m not sure if it fits here and how to use it in this case.
Thank you!