AsyncCheckpointIO — PyTorch Lightning 2.5.1.post0 documentation (original) (raw)

class lightning.pytorch.plugins.io.AsyncCheckpointIO(checkpoint_io=None)[source]¶

Bases: _WrappingCheckpointIO

AsyncCheckpointIO enables saving the checkpoints asynchronously in a thread.

Parameters:

checkpoint_io¶ (Optional[CheckpointIO]) – A checkpoint IO plugin that is used as the basis for async checkpointing.

save_checkpoint(*args, **kwargs)[source]¶

Uses the ThreadPoolExecutor to save the checkpoints using the base checkpoint_io.

Return type:

None

teardown()[source]¶

This method is called to close the threads.

Return type:

None