Change register_backward_hook to register_full_backward_hook

Summary: register_backward_hook is getting deprecated so we merge it to register_full_backward_hook Differential Revision: D48725761
HuanyuZhang · Jun 10, 2024 · 4440412 · 4440412
1 parent 202c58a
commit 4440412
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/opacus/grad_sample/grad_sample_module.py b/opacus/grad_sample/grad_sample_module.py
@@ -207,7 +207,7 @@ def add_hooks(
             )
 
             self.autograd_grad_sample_hooks.append(
-                module.register_backward_hook(
+                module.register_full_backward_hook(
                     partial(
                         self.capture_backprops_hook,
                         loss_reduction=loss_reduction,