Update optimizer.py

torch/optim/optimizer.py

@@ -645,11 +645,11 @@ def register_load_state_dict_pre_hook(
         hook: Callable[["Optimizer", StateDict], Optional[StateDict]],
         prepend: bool = False,
     ) -> RemovableHandle:
-        r"""Register a state dict pre-hook which will be called before
-        :meth:`~torch.optim.Optimizer.state_dict` is called. It should have the
+        r"""Register a load_state_dict pre-hook which will be called before
+        :meth:`~torch.optim.Optimizer.load_state_dict` is called. It should have the
         following signature::
 
-            hook(optimizer) -> None
+            hook(optimizer, state_dict) -> state_dict or None
 
         The ``optimizer`` argument is the optimizer instance being used and the
         ``state_dict`` argument is a shallow copy of the ``state_dict`` the user
@@ -683,11 +683,11 @@ def register_load_state_dict_pre_hook(
     def register_load_state_dict_post_hook(
         self, hook: Callable[["Optimizer"], None], prepend: bool = False
     ) -> RemovableHandle:
-        r"""Register a state dict post-hook which will be called after
-        :meth:`~torch.optim.Optimizer.state_dict` is called. It should have the
+        r"""Register a load_state_dict post-hook which will be called after
+        :meth:`~torch.optim.Optimizer.load_state_dict` is called. It should have the
         following signature::
 
-            hook(optimizer, state_dict) -> state_dict or None
+            hook(optimizer) -> None
 
         The ``optimizer`` argument is the optimizer instance being used.