Update on "Allow Module forward-pre and forward hooks to take kwargs"

mrshenli · mrshenli · commit d940db00e977 · 2022-11-22T22:32:49.000Z
closes #35643 This PR is mostly borrowed from #82042. Thanks Padarn for implementing the first version and debugging into the errors. Based on the discussion in #82042 this PR adds a with_kwargs argument to register_forward_pre_hook and register_forward_hook methods. When the arg is set to true, the provided hook must accept kwargs args. Under the hook, this PR adds a `_forward_pre_hooks_with_kwargs` and a `_forward_hook_with_kwargs` set to keep track of which hooks accept kwargs. Differential Revision: [D41431111](https://our.internmc.facebook.com/intern/diff/D41431111) cc albanD mruberry jbschlosser walterddr kshitij12345 saketh-are [ghstack-poisoned]
diff --git a/torch/nn/modules/module.py b/torch/nn/modules/module.py
@@ -1495,11 +1495,13 @@ def _call_impl(self, *args, **kwargs):
                 if hook_id in self._forward_pre_hooks_with_kwargs:
                     result = hook(self, args, kwargs)  # type: ignore[misc]
                     if result is not None:
-                        assert isinstance(result, tuple) and len(result) == 2, (
-                            "forward pre-hook must return None or a tuple of "
-                            f"(new_args, new_kwargs), but got {result}."
-                        )
-                        args, kwargs = result
+                        if isinstance(result, tuple) and len(result) == 2:
+                            args, kwargs = result
+                        else:
+                            raise RuntimeError(
+                                "forward pre-hook must return None or a tuple "
+                                f"of (new_args, new_kwargs), but got {result}."
+                            )
                 else:
                     result = hook(self, args)
                     if result is not None:
diff --git a/torch/utils/hooks.py b/torch/utils/hooks.py
@@ -37,7 +37,7 @@ def remove(self) -> None:
 
         if self.extra_dict_ref is not None:
             extra_dict = self.extra_dict_ref()
-            if extra_dict is not None:
+            if extra_dict is not None and self.id in extra_dict:
                 del extra_dict[self.id]
 
     def __getstate__(self):
