Update hooks.py (#1564)

k2-fsa · Mar 20, 2024 · d5cd78a · d5cd78a
1 parent 9bd3085
commit d5cd78a
Showing 1 changed file with 10 additions and 2 deletions.
diff --git 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():