Solved by ptrblck in post #4 
              
              
                No, you wouldn’t need a for loop and for batched input this should work: 
x[torch.arange(x.size(0)).unsqueeze(1), torch.arange(x.size(1)), i]