How did you define idx? Does it store the “flat” indices, i.e. [0, 256*512]?
Could you try this:
x = torch.randn(1, 1, 256, 512)
idx = torch.arange(256*512).long()
idx = idx[torch.randperm(256*512)]

y = x.view(-1)[idx]

Let me know, if this suits your use case.