github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-05-10 09:27:50 +00:00
Code Issues Packages Projects Releases Wiki Activity

Fix docstr for zero statedict (#3185)

Browse Source
This commit is contained in:
YH 2023-03-21 12:48:21 +09:00 committed by GitHub
parent 7bc0afc901
commit 9d644ff09f
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 1 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
colossalai/zero/sharded_model/sharded_model_v2.py
View File

@ -494,6 +494,7 @@ class ShardedModelV2(nn.Module):
error_msgs (list of str): error messages should be added to this error_msgs (list of str): error messages should be added to this
list, and will be reported together in list, and will be reported together in
:meth:`~torch.nn.Module.load_state_dict` :meth:`~torch.nn.Module.load_state_dict`
shard_strategy (Optional[BaseShardStrategy], optional): A shard strategy to manage shard behavior. Defaults to None.
""" """
for hook in self._load_state_dict_pre_hooks.values(): for hook in self._load_state_dict_pre_hooks.values():
hook(state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs) hook(state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs)