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

[jjhhan]

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?

[cucdengjunli]

same question