Hi all:
Suppose I have a 3D tensor: x , shaped 2 x 247 x 4 and a 1D index tensor of length 2: y = [a, b], whose length corresponds to the size of the 1st dimension of the 3D tensor. What I want to do is index out x[0, a, :] and x[1, b, :], and have my output be shape 2 x 4, how would I do that? And of course in general, I want this to apply to 3D tensor of shape batch_size x 247 x 4, and my 1D index of length = batch_size. Please share your thoughts! Thank you!
Example:
I don’t quite understand your use case.
Let’s assume you have a tensor x = torch.randn(2, 247, 4).
Now you would like to create an index tensor y containing two values? What do you mean by “whole length corresponds to the size of the 1st dim”?
Do you have two values in y or is it of shape [2, 247]?
If you have two values and index x with y, you would get a tensor of shape [2, 2, 4].
Could you share some pseudo code so that I could have a look and understand the use case a bit better?
Thanks for the reply! Sorry I don’t know if I can format and include code directly into my post, but I added an example using screenshot.
