Solved by agadetsky in post #3 
              
              
                Tensorized version can be implemented as follows: 
def smart_sort(x, permutation):
    d1, d2 = x.size()
    ret = x[
        torch.arange(d1).unsqueeze(1).repeat((1, d2)).flatten(),
        permutation.flatten()
    ].view(d1, d2)
    return ret

And it sorts as you want: 
a = torch.tensor([[1, 0, …