Solved by ptrblck in post #2 
              
              
                Wouldn’t the result be [0, 4] in this case? 
If so, this should work: 
a = torch.tensor(([0,1,2],[3,4,5]))
idx = torch.tensor([0, 1])
a[torch.arange(a.size(0)), idx]
# > tensor([0, 4])