comment

AlibabaPAI · Sep 4, 2024 · a729b08 · a729b08
1 parent 8dfe87e
commit a729b08
Showing 1 changed file with 6 additions and 1 deletion.
diff --git a/torchacc/dist/fsdp.py b/torchacc/dist/fsdp.py
@@ -247,7 +247,7 @@ def optim_state_dict(
 
         Returns:
             Dict[str, Any]: A :class:`dict` containing the optimizer state for
-            ``model``. The sharding of the optimizer state is based on
+            self.model. The sharding of the optimizer state is based on
             ``state_dict_type``.
         """
         # we only support FULL_STATE_DICT and flatten parameters now
@@ -331,6 +331,11 @@ def load_optim_state_dict(
                 The default value is FULL_STATE_DICT.
             rank0_only: (bool = True): control whether load state_dict only from
                 rank0 at the begining.
+        
+        Returns:
+            Dict[str, Any]: A :class:`dict` containing the optimizer state for
+            self.model. The sharding of the optimizer state is based on
+            ``state_dict_type``.
         """
         # we only support FULL_STATE_DICT and flatten parameters now
         if state_dict_type != StateDictType.FULL_STATE_DICT: