Skip to content

Commit aa35583

Browse files
committed
add docstring
1 parent 2dcafd0 commit aa35583

File tree

3 files changed

+13
-2
lines changed

3 files changed

+13
-2
lines changed

pytorch_lightning/accelerators/accelerator.py

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -481,6 +481,12 @@ def connect_precision_plugin(self, plugin: PrecisionPlugin) -> None:
481481
self.setup_precision_plugin(plugin)
482482

483483
def save_checkpoint(self, checkpoint: Dict[str, Any], filepath) -> None:
484+
"""Save model/training states as a checkpoint file through state-dump and file-write.
485+
486+
Args:
487+
checkpoint: dict containing model and trainer state
488+
filepath: write-target file's path
489+
"""
484490
self.training_type_plugin.save_checkpoint(checkpoint, filepath)
485491

486492
@property

pytorch_lightning/plugins/training_type/tpu_spawn.py

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -300,9 +300,8 @@ def save_checkpoint(self, checkpoint: Dict[str, Any], filepath: str) -> None:
300300
"""Save model/training states as a checkpoint file through state-dump and file-write.
301301
302302
Args:
303-
trainer: PyTorch Lightning Trainer
303+
checkpoint: dict containing model and trainer state
304304
filepath: write-target file's path
305-
weights_only: saving model weights only
306305
"""
307306
# Todo: TypeError: 'mappingproxy' object does not support item assignment
308307
self.save({k: v for k, v in checkpoint.items() if k != "callbacks"}, filepath)

pytorch_lightning/plugins/training_type/training_type_plugin.py

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -198,6 +198,12 @@ def setup_optimizers_in_pre_dispatch(self) -> bool:
198198
return False
199199

200200
def save_checkpoint(self, checkpoint: Dict[str, Any], filepath: str) -> None:
201+
"""Save model/training states as a checkpoint file through state-dump and file-write.
202+
203+
Args:
204+
checkpoint: dict containing model and trainer state
205+
filepath: write-target file's path
206+
"""
201207
# dump states as a checkpoint dictionary object
202208
if self.is_global_zero:
203209
checkpoint = self.on_save(checkpoint)

0 commit comments

Comments
 (0)