Hook for Fully Formed Checkpoints

[dcharatan]

I want to create a hook that uploads checkpoints to cloud storage (e.g. AWS, Azure). I tried using the on_save_checkpoint hook as follows:

def on_save_checkpoint(self, trainer: pl.Trainer, pl_module: pl.LightningModule, checkpoint: Dict[str, Any]) -> dict:
    checkpoint_bytes = io.BytesIO()
    torch.save(checkpoint, checkpoint_bytes)
    # Upload the BytesIO...

However, states for optimizers, learning rate schedulers, etc. are added to the checkpoint dict after on_save_checkpoint is called. Is there an elegant way to create a hook that receives fully formed checkpoints?

edit: sorry, this is a duplicate -- GitHub was giving me errors when I posted

[akihironitta]

See #11704.