When you call save_checkpoint, is the model var a DDP instance? If yes, you might need to save model.module instead? But I don’t think that’s the reason for the error jump. When you load the module, if you do not use DDP or DP (just a local model), is the loss after recovery as expected?
I might miss sth, looks like in the “resume model” part, the model state is loaded to CPU and not moved to local_rank before passing to DDP ctors? Or is the model already on the correct device before loading state dict?