From 2102038e0e3e1989f4fb4d621c43c43ac1ca23bc Mon Sep 17 00:00:00 2001
From: Daniel Povey <dpovey@gmail.com>
Date: Thu, 1 Dec 2022 16:23:50 +0800
Subject: [PATCH] Fix bug in diagnostics.py

---
 icefall/diagnostics.py | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/icefall/diagnostics.py b/icefall/diagnostics.py
index 2a937c6b8..e596c0028 100644
--- a/icefall/diagnostics.py
+++ b/icefall/diagnostics.py
@@ -543,6 +543,10 @@ def attach_diagnostics(
                     _model_diagnostic[f"{_name}.grad[{i}]"].accumulate(o,
                                                                        class_name=type(_module).__name__)
 
+
+        module.register_forward_hook(forward_hook)
+        module.register_backward_hook(backward_hook)
+
         if type(module).__name__ in  ["Sigmoid", "Tanh", "ReLU", "TanSwish", "Swish", "DoubleSwish"]:
             # For these specific module types, accumulate some additional diagnostics
             # that can help us improve the activation function.  These require a lot of memory,