From 2da0cd183fedb7a05a7ada03fdd41d27e979cfcd Mon Sep 17 00:00:00 2001 From: yhna940 Date: Tue, 21 Mar 2023 00:09:14 +0900 Subject: [PATCH] Fix docstr for zero statedict --- 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 094f7d76a86d..12e8f65d4a35 100644 --- a/colossalai/zero/sharded_model/sharded_model_v2.py +++ b/colossalai/zero/sharded_model/sharded_model_v2.py @@ -494,6 +494,7 @@ def _colo_load_from_state_dict(self, 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)