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, 2, 2, 1], [0, 2, 1, 2, 0]])
b = torch.tensor([[3, 2, 4, 1, 0], [2, 1, 4, 3, 3]])
smart_sort(a, b)

output: tensor([[2, 2, 1, 0, 1],
               [1, 2, 0, 2, 2]])