scvi.model.base.PyroJitGuideWarmup.on_save_checkpoint

PyroJitGuideWarmup.on_save_checkpoint(trainer, pl_module, checkpoint)

Called when saving a model checkpoint, use to persist state.

Parameters
trainer : TrainerTrainer

the current Trainer instance.

pl_module : LightningModuleLightningModule

the current LightningModule instance.

checkpoint : {str: Any}Dict[str, Any]

the checkpoint dictionary that will be saved.

Return type

dictdict

Returns

The callback state.