Slicing via lists doesn’t mess up any gradient computations. Explicitly:
import torch
from torch.nn.utils.rnn import pad_sequence

h1 = torch.tensor([1., 1, 2, 3], requires_grad=True)
h2 = torch.tensor([5., 5], requires_grad=True)
h = pad_sequence([h1, h2])
lengths = torch.tensor([len(h1), len(h2)]).long()
batch_size = len(lengths)

# list slicing
final_states_list = h[[l-1 for l in lengths], [i for i in range(batch_size)]]

# gather via mask
mask = (lengths-1).view(1, -1).expand(h.size(0), h.size(1))
final_states_mask = h.gather(0, mask)[0]

#check grads both ways
final_states_list.sum().backward(retain_graph=True)
print(h1.grad, h2.grad)
# tensor([0., 0., 0., 1.]) tensor([0., 1.])
h1.grad.zero_()
h2.grad.zero_()
final_states_mask.sum().backward()
print(h1.grad, h2.grad)
# tensor([0., 0., 0., 1.]) tensor([0., 1.])

I’m not sure what the advantage of using gather() here is, since all rows but the first of the gather()'d tensor are thrown away. The list slicing seems cleaner.