Hello,
As far as I know, there is no such function to get a 2D tensor sorted from another 2D tensor. However, with a simple for loop and indexing, it is possible to achieve the desired result:
    a = torch.tensor([[1, 0, 2, 2, 1], [0, 2, 1, 2, 0]], dtype=torch.int32)
    # index tensor (b) must be of type long, byte or bool
    b = torch.tensor([[3, 2, 4, 1, 0], [2, 1, 4, 3, 3]], dtype=torch.long)
    for i in range(b.shape[0]):
        # change the order of the ith tensor of a with the indexes of b[i]
        c = a[i][b[i]]
        print(c)
    # tensor([2, 2, 1, 0, 1], dtype=torch.int32)
    # tensor([1, 2, 0, 2, 2], dtype=torch.int32)

Don’t know if this helps!