For every batch element of output of GRU I’m taking a certain timestep slice that corresponds to the last timestep for that sequence. This discards hidden states corresponding to padding time steps - that’s what I mean by “undoing the padding”
Thanks for your suggestion! In my original post notation it should be the following, right?
def index_dim(tensor, index):
  return tensor[torch.arange(len(tensor)), index]