I’m not exactly sure which issue you are currently trying to solve.
To save and restore the training you would need to save all trained parameters, buffers, etc.
To do so, you would need to save the state_dict from all models, the optimizer, and schedulers.
This will make sure that the internal “states” of these objects are properly restored and you could continue the training or perform the inference.
On the other hand, if you are trying to debug the issue that your models are not training properly, make sure the computation graph is not detached e.g. by checking the .grad attributes of all parameters which should be trained.



 KURUVILLA_ABRAHAM:

I have few transformations done after getting the model output and cycle gan like getting color information and then using it as a criterion.


If these transformations do not have an internal state you wouldn’t have to save anything (there won’t be anything to save besides the actual source code).