Browse Source

Fix docstr for zero statedict (#3185)

pull/3191/head
YH 2 years ago committed by GitHub
parent
commit
9d644ff09f
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
  1. 1
      colossalai/zero/sharded_model/sharded_model_v2.py

1
colossalai/zero/sharded_model/sharded_model_v2.py

@ -494,6 +494,7 @@ class ShardedModelV2(nn.Module):
error_msgs (list of str): error messages should be added to this
list, and will be reported together in
: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():
hook(state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs)

Loading…
Cancel
Save