Solved by smth in post #2 
              
              
                it should be: 
idx_permute = torch.from_numpy(rng_permute.permutation(784), dtype=torch.int64)
transform = torchvision.transforms.Compose([torchvision.transforms.ToTensor(),
              torchvision.transforms.Lambda(lambda x: x.view(-1)[idx_permute].view(1, 28, 28) )])