No idea if this is the most efficient method but I think it gets the job done.
Change your line:
last_seq_items = torch.index_select(output, 1, last_seq_idxs) 

to this:
last_seq_items = output[range(output.shape[0]), last_seq_idxs, :]

By using range(), we are able to give it a different index to take at each 0th dimension.