Get it wrong once, never recover. 
But here is what I actually have tested:
data = torch.randn(64,20,1024)
last_elements = torch.randint(0, 20, (64,)).long()
myrange = torch.arange(0,64, dtype=torch.int64)
print(data[myrange, last_elements])

Best regards
Thomas