Lightning is freezing

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