Skip to content
Discussion options

You must be logged in to vote

you can replace the model state_dict inside the checkpoint

class LitModel(LightningModule):
    ...
    
    def on_save_checkpoint(self, checkpoint):
        with ema.average_parameters():
            checkpoint['state_dict'] = self.state_dict()

Replies: 1 comment

Comment options

You must be logged in to vote
0 replies
Answer selected by maxmatical
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment