Fix bug where adding a hook could replace an existing hook.

We were keying hooks by RemovableHandle id. However, we don't hold onto
handles and ids of dead objects can be reused. This replaces id(handle)
with a global counter.

Author: colesbury

torch/autograd/function.py

@@ -111,7 +111,7 @@ def _register_hook(backward_hooks, hook):
         if backward_hooks is None:
             backward_hooks = OrderedDict()
         handle = hooks.RemovableHandle(backward_hooks)
-        backward_hooks[id(handle)] = hook
+        backward_hooks[handle.id] = hook
         return backward_hooks, handle
 
     def forward(self, *input):

torch/autograd/variable.py

@@ -180,7 +180,7 @@ def register_hook(self, hook):
             if self.creator is not None:
                 self.creator._register_hook_dict(self)
         handle = hooks.RemovableHandle(self._backward_hooks)
-        self._backward_hooks[id(handle)] = hook
+        self._backward_hooks[handle.id] = hook
         return handle
 
     def reinforce(self, reward):

torch/nn/modules/module.py

@@ -179,7 +179,7 @@ def register_backward_hook(self, hook):
         that removes the hook from the module.
         """
         handle = hooks.RemovableHandle(self._backward_hooks)
-        self._backward_hooks[id(handle)] = hook
+        self._backward_hooks[handle.id] = hook
         return handle
 
     def register_forward_hook(self, hook):
@@ -195,7 +195,7 @@ def register_forward_hook(self, hook):
         that removes the hook from the module.
         """
         handle = hooks.RemovableHandle(self._forward_hooks)
-        self._forward_hooks[id(handle)] = hook
+        self._forward_hooks[handle.id] = hook
         return handle
 
     def __call__(self, *input, **kwargs):

torch/utils/hooks.py

@@ -4,14 +4,17 @@
 class RemovableHandle(object):
     """A handle which provides the capability to remove a hook."""
 
+    next_id = 0
+
     def __init__(self, hooks_dict):
         self.hooks_dict_ref = weakref.ref(hooks_dict)
+        self.id = RemovableHandle.next_id
+        RemovableHandle.next_id += 1
 
     def remove(self):
         hooks_dict = self.hooks_dict_ref()
-        key = id(self)
-        if hooks_dict is not None and key in hooks_dict:
-            del hooks_dict[key]
+        if hooks_dict is not None and self.id in hooks_dict:
+            del hooks_dict[self.id]
 
     def __enter__(self):
         return self