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?