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

Question

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 1

You can use the function torch.randperm to get a set of indices that act as a random permutation. The following is a small example of getting a random permutation, then applying it to both the a and b tensors:

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