img = torch.arange(3*32*32, dtype=torch.float32).reshape(3, 32, 32)
sh  = shuffle_rows(img)

print(sh.shape)                                     # torch.Size([3, 32, 32])

# every value still present, exactly once
print(torch.equal(sh.flatten().sort().values,
                  img.flatten().sort().values))     # True

# and it's perfectly reversible
inv = torch.empty_like(row_perm)
inv[row_perm] = torch.arange(32)
print(torch.equal(sh[:, inv, :], img))              # True