From 9d644ff09f2b044c984328e08357a68d98ab17f3 Mon Sep 17 00:00:00 2001 From: YH <100389977+yhna940@users.noreply.github.com> Date: Tue, 21 Mar 2023 12:48:21 +0900 Subject: [PATCH] Fix docstr for zero statedict (#3185) --- colossalai/zero/sharded_model/sharded_model_v2.py | 1 + 1 file changed, 1 insertion(+) diff --git a/colossalai/zero/sharded_model/sharded_model_v2.py b/colossalai/zero/sharded_model/sharded_model_v2.py index 094f7d76a..12e8f65d4 100644 --- a/colossalai/zero/sharded_model/sharded_model_v2.py +++ b/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)