diff --git a/egs/librispeech/ASR/pruned_transducer_stateless4b/train.py b/egs/librispeech/ASR/pruned_transducer_stateless4b/train.py
index f35e4e08f..ba88060e3 100755
--- a/egs/librispeech/ASR/pruned_transducer_stateless4b/train.py
+++ b/egs/librispeech/ASR/pruned_transducer_stateless4b/train.py
@@ -703,7 +703,7 @@ def train_one_epoch(
 
         if params.batch_idx_train % 2000 == 0 and params.batch_idx_train > 0:
             mmodel = model.module if hasattr(model, 'module') else model
-            mmodel.encoder.diagonalize()
+            mmodel.diagonalize()
             #optimizer.reset()
 
         params.batch_idx_train += 1