Update hook registering to newest API

szymonmaszke · Jun 13, 2022 · 1eff7c5 · 1eff7c5
1 parent 24669cd
commit 1eff7c5
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/torchlayers/regularization.py b/torchlayers/regularization.py
@@ -161,7 +161,7 @@ def __init__(self, module, weight_decay, name: str = None):
         self.weight_decay = weight_decay
         self.name = name
 
-        self.hook = self.module.register_backward_hook(self._weight_decay_hook)
+        self.hook = self.module.register_full_backward_hook(self._weight_decay_hook)
 
     def remove(self):
         self.hook.remove()