Replace register_backward_hook with register_full_backward_hook (#720)

iden-kalemaj · facebook-github-bot · commit 8cbf8e0733ea · 2025-04-04T10:44:13.000-07:00
Summary: Pull Request resolved: #720 register_backward_hook is deprecated and may lead to errors in gradient calculation. We switch to the supported register_full_backward_hook. Reviewed By: HuanyuZhang Differential Revision: D68562558 fbshipit-source-id: 58bf997a617da6eb097c6b159ed9354d4c8320a2
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,

Original file line number	Diff line number	Diff line change
`@@ -207,7 +207,7 @@ def add_hooks(`
`207`	`207`	`)`
`208`	`208`
`209`	`209`	`self.autograd_grad_sample_hooks.append(`
`210`		`- module.register_backward_hook(`
	`210`	`+ module.register_full_backward_hook(`
`211`	`211`	`partial(`
`212`	`212`	`self.capture_backprops_hook,`
`213`	`213`	`loss_reduction=loss_reduction,`