Solved by tumble-weed in post #2 
              
              
                you can do tensor.view(-1) to view it as a flat tensor, apply the shuffling and then view it in the original shape again