Solved by KarthikR in post #2 
              
              
                I tried with a smaller tensor but should work for you: 
input1  = torch.randn(2,2,3) 
input2  = torch.randn(2,2,3) 
print(input2) 
values1, indices1 = input1.max(dim=-1,  keepdim=True) 
print(’\n’,indices1, ‘\n’) 
print(torch.gather(input2, -1, indices1))