Unable to backward when training

[VietPT3502]

Hello, i have try to train using your code. And when loss.backward(), this take forever to run a single iter. When i tried to use another basic model, for example Unet the problems is fixed. So is there any problems happening in your model's code and can you give me some advice. Thanks for advance!