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:

Screen Shot 2018-10-22 at 9.23.21 AM.png2034×772 111 KB