Is there a loopless native PyTorch equivalent of this op?
def index_dim(tensor, dim, index):
  return torch.stack([tensor[k].select(dim - 1, i) for k, i in enumerate(index.tolist())])

t = torch.rand(2, 4, 3)
i = torch.LongTensor([0, 1])
print(index_dim(t, 1, i).shape)
# torch.Size([2, 3])


In practice I need this to undo padding from a GRU output (I know I could use PackedSequence format, but for now I’d like to try padded format as it’s more compatible with ConvNet models) and get an embedding from the last meaningful timestep for every sequence in the batch.