github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-09-21 17:40:33 +00:00
Code Issues Packages Projects Releases Wiki Activity

polish utils docstring (#620)

Browse Source
This commit is contained in:
ver217
2022-04-01 16:36:47 +08:00
committed by GitHub
parent e619a651fb
commit 369a288bf3
4 changed files with 42 additions and 48 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
colossalai/utils/checkpointing.py
View File

@@ -175,7 +175,7 @@ def load_checkpoint(checkpoint_path: str,
If strict is True, then the keys of state_dict must exactly match the keys returned
by this module’s state_dict() function.
Args:
Args:
checkpoint_path (str): The exact and matched checkpoint_path directory to retrieve appropriate state_dict.
model (:class:`torch.nn.Module`): Model to reload parameters and buffers.
optimizer (Union[:class:`torch.optim.Optimizer`, :class:`colossalai.nn.optimizer`]): Optimizer to recuperate.