When should the model be saved, and how should it be implemented?

[ZH-UCAS]

When should the model be saved, and how should it be implemented? I use torch.save(self.unet.state_dict(), f"{self.savepath}unet_checkpoint_{timestamp}epoch{epoch}.pth") and torch.save(self.diffusion.state_dict(), f"{self.savepath}diffusion_checkpoint_{timestamp}epoch{epoch}.pth") to separately store the weights of the unet and diffusion models. I save them after the backward pass and optimizer adjustments. However, when loading the unet, some parameters seem to be missing, and the generated samples after loading do not match the quality during the training phase. As a beginner, please advise on when and how I should save and load the model.