I have a tensor of indices > 0 and < 256, of shape [1,5,8], and a tensor of images of shape [1,5,256,3,84,84]. All I want to do is create a new tensor with the images at those indices - i.e. a tensor of size [1,5,8,3,84,84]. But nothing I do seems to work. torch.index_select seems to do this across one axis, but not multiple, and torch.gather would work, except it requires the index tensor and input tensor to be the same size. Do I really have to resort to a for loop here, or is there a nice way to do this?
torch.gather should do just this
You don’t need them to have the same shape, or you can just I squeeze it to the proper dimension
Except torch.gather does require the same number of dimensions. I assume that if I take the index tensor to size [1,5,8,1,1,1], it’ll just give me the first pixel of each image, so how do I leave the images intact when taking the parts?