Merge pull request #27 from tfjgeorge/full_backward_hooks

switches to full backward hooks in jacobian generator
tfjgeorge · Jul 5, 2021 · 522a9a4 · 522a9a4
2 parents d6cbe15 + 4a3f621
commit 522a9a4
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/nngeometry/generator/jacobian.py b/nngeometry/generator/jacobian.py
@@ -630,7 +630,7 @@ def _add_hooks(self, hook_x, hook_gy, mods):
         handles = []
         for m in mods:
             handles.append(m.register_forward_pre_hook(hook_x))
-            handles.append(m.register_backward_hook(hook_gy))
+            handles.append(m.register_full_backward_hook(hook_gy))
         return handles
 
     def _hook_savex(self, mod, i):
@@ -1117,4 +1117,4 @@ def _get_dataloader(self, examples):
             return examples
         else:
             return DataLoader(TensorDataset(*examples),
-                              batch_size=len(examples[0]))
+                              batch_size=len(examples[0]))