Solved by ptrblck in post #2 
              
              
                This should work: 
a = torch.randn(2, 5, 2)
indices = a[:, :, 0].sort()[1]
a_sorted = a[torch.arange(a.size(0)).unsqueeze(1), indices]

print(a)
print(a_sorted)