Solved by fmassa in post #2 
              
              
                torch.max returns both the max values as well as the indices. 
so you can do 
values, indices = tensor.max(0)
values, indices = torch.max(tensor, 0)