From d5cd78a63722f2c105d0468f0a51f3ac8f404ad8 Mon Sep 17 00:00:00 2001
From: zr_jin <peter.jin.cn@gmail.com>
Date: Wed, 20 Mar 2024 16:43:45 +0800
Subject: [PATCH] Update hooks.py (#1564)

---
 icefall/hooks.py | 12 ++++++++++--
 1 file changed, 10 insertions(+), 2 deletions(-)

diff --git a/icefall/hooks.py b/icefall/hooks.py
index 398a5f689d..1c5bd2ae68 100644
--- a/icefall/hooks.py
+++ b/icefall/hooks.py
@@ -1,4 +1,6 @@
-# Copyright  2021-2022  Xiaomi Corporation  (authors: Zengwei Yao, Daniel Povey)
+# Copyright  2021-2024  Xiaomi Corporation  (authors: Zengwei Yao,
+#                                                     Daniel Povey,
+#                                                     Zengrui Jin,)
 #
 # See ../../LICENSE for clarification regarding multiple authors
 #
@@ -77,7 +79,13 @@ def param_backward_hook(grad, _name=name):
             if not torch.isfinite(grad.to(torch.float32).sum()):
                 logging.warning(f"The sum of {_name}.param_grad is not finite")
 
-        parameter.register_hook(param_backward_hook)
+        try:
+            parameter.register_hook(param_backward_hook)
+        except:
+            logging.warning(
+                f"Warning: could not register backward hook for parameter {name}, "
+                f"it might not be differentiable."
+            )
 
 
 def _test_inf_check_hooks():