From 8ccd061051292d783d53a02b1648860d68c8f146 Mon Sep 17 00:00:00 2001
From: Daniel Povey <dpovey@gmail.com>
Date: Sat, 11 Feb 2023 17:31:21 +0800
Subject: [PATCH] Fix bug where attn_mask was not passed in.

---
 egs/librispeech/ASR/pruned_transducer_stateless7/zipformer.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/egs/librispeech/ASR/pruned_transducer_stateless7/zipformer.py b/egs/librispeech/ASR/pruned_transducer_stateless7/zipformer.py
index 80797c681..4376967a0 100644
--- a/egs/librispeech/ASR/pruned_transducer_stateless7/zipformer.py
+++ b/egs/librispeech/ASR/pruned_transducer_stateless7/zipformer.py
@@ -479,9 +479,9 @@ class Zipformer(EncoderInterface):
 
         attn_mask = torch.logical_or(src_c > tgt_c,
                                      src_c < tgt_c - left_context_chunks)
-        if __name__ == "__main__":
+        if __name__ == "__main__" or random.random() < 0.1:
             logging.info(f"attn_mask = {attn_mask}")
-
+        return attn_mask
 
 
 def _whitening_schedule(x: float, ratio: float = 2.0) -> ScheduledFloat: