Hook for Fully Formed Checkpoints
I would like to create a hook that automatically uploads checkpoints to the cloud (e.g., AWS, Azure) when they're created. I tried using on_save_checkpoint roughly like this:
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 somewhere...

However, states for optimizers, schedulers, AMP, etc. are added after on_save_checkpoint hooks are called. Is there an elegant way to create a hook that receives the fully formed checkpoint state?
hey @dcharatan !
I'd rather suggest using the remote filesystems. You can also specify the remote path inside ModelCheckpoint.
or use CheckpointIO plugin.