Freezing on the manual backward scenario

[abdullahbas]

Hi all,

I am trying to use manual backward because I split my input images into 2D slices from 3D images. If I return loss from the end of training_step it freezes at the last batch of the first epoch if I don’t return any loss still the same. If I return loss inside the for loop then it works but with only one part of my batch which is not appropriate.

Thanks
Abdullah

[awaelchli]

Hi, can you post a small running example? I don't exactly understand what you have tried to do.