If I understand your use case correctly, you would like to be able to revert the shuffling?
If so, this should work:
# setup
N, M, K = 2, 5, 2
x = torch.arange(N*M*K).view(N, M, K)
print(x)
> tensor([[[ 0,  1],
           [ 2,  3],
           [ 4,  5],
           [ 6,  7],
           [ 8,  9]],

          [[10, 11],
           [12, 13],
           [14, 15],
           [16, 17],
           [18, 19]]])

# shuffle
idx = torch.randperm(x.size(1))
y = x[:, idx]
print(y)
> tensor([[[ 8,  9],
           [ 2,  3],
           [ 0,  1],
           [ 4,  5],
           [ 6,  7]],

          [[18, 19],
           [12, 13],
           [10, 11],
           [14, 15],
           [16, 17]]])

# reverse the shuffling
idx_inv = torch.sort(idx).indices
print(y[:, idx_inv])
> tensor([[[ 0,  1],
           [ 2,  3],
           [ 4,  5],
           [ 6,  7],
           [ 8,  9]],

          [[10, 11],
           [12, 13],
           [14, 15],
           [16, 17],
           [18, 19]]])