Solved by ptrblck in post #2 
              
              
                The following indexing should work: 
x = torch.randn(16, 1580, 201)
idx = torch.tensor(
    [1580, 959, 896, 881, 881, 881, 881, 881, 881, 881, 881, 881, 881, 335, 254, 219]
)
idx = idx - 1  # 0-based index
y = x[torch.arange(x.size(0)), idx]