use full backward hook

justanhduc · Dec 6, 2021 · c629664 · c629664
1 parent 2b24a3f
commit c629664
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/neural_monitor/monitor.py b/neural_monitor/monitor.py
@@ -120,7 +120,7 @@ def _forward_hook(module, input, output):
             def _backward_hook(module, grad_input, grad_output):
                 _TRACKS['grad_' + name + '_output'] = tuple([grad_out.detach() for grad_out in grad_output])
 
-            hooks['grad_' + name + '_output'] = x.register_backward_hook(_backward_hook)
+            hooks['grad_' + name + '_output'] = x.register_full_backward_hook(_backward_hook)
     else:
         if direction in ('forward', 'all', None):
             _TRACKS[name] = x.detach()