Solved by ptrblck in post #2 
              
              
                Assuming you would like to get: 
x[0, 903]
x[1, 815]
...

then this should work: 
x = torch.randn(5, 1000)
idx = torch.tensor([903, 815, 214, 282, 694])

x[torch.arange(x.size(0)), idx]