Shuffling two 2D tensors in PyTorch and maintaining same order correlation

Is it possible to shuffle two 2D tensors in PyTorch by their rows, but maintain the same order for both? I know you can shuffle a 2D tensor by rows with the following code:

a=a[torch.randperm(a.size()[0])]

To elaborate: If I had 2 tensors

a = torch.tensor([[1, 1, 1, 1, 1],
            [2, 2, 2, 2, 2],
            [3, 3, 3, 3, 3]])

b = torch.tensor([[4, 4, 4, 4, 4],
            [5, 5, 5, 5, 5],
            [6, 6, 6, 6, 6]])

And ran them through some function/block of code to shuffle randomly but maintain correlation and produce something like the following

a = torch.tensor([[2, 2, 2, 2, 2],
            [1, 1, 1, 1, 1],
            [3, 3, 3, 3, 3]])

b = torch.tensor([[5, 5, 5, 5, 5],
            [4, 4, 4, 4, 4],
            [6, 6, 6, 6, 6]])

My current solution is converting to a list, using the random.shuffle() function like below.

a_list = a.tolist()
b_list = b.tolist()
temp_list = list(zip(a_list , b_list ))
random.shuffle(temp_list) # Shuffle
a_temp, b_temp = zip(*temp_list)
a_list, b_list = list(a_temp), list(b_temp)
            
# Convert back to tensors
a = torch.tensor(a_list)
b = torch.tensor(b_list)

This takes quite a while and was wondering if there is a better way.

Answer

You mean

indices = torch.randperm(a.size()[0])
a=a[indices]
b=b[indices]

?