how can i modify the code to train on 3D data like x = torch.randn(3, 3, 64, 64, 64)

[JJH]

Thanks for your great job. I want to test the code on the 3D data， Could you please give me some advice on how to change the code?

[Dylan_邓珺礼]

same question