diff --git a/egs/librispeech/ASR/pruned_transducer_stateless2/conformer.py b/egs/librispeech/ASR/pruned_transducer_stateless2/conformer.py
index bcd419fb7..ab46e233b 100644
--- a/egs/librispeech/ASR/pruned_transducer_stateless2/conformer.py
+++ b/egs/librispeech/ASR/pruned_transducer_stateless2/conformer.py
@@ -158,7 +158,7 @@ class Conformer(EncoderInterface):
         if not is_jit_tracing():
             assert x.size(0) == lengths.max().item()
 
-        src_key_padding_mask = make_pad_mask(lengths)
+        src_key_padding_mask = make_pad_mask(lengths, x.size(0))
 
         if self.dynamic_chunk_training:
             assert (
diff --git a/egs/librispeech/ASR/pruned_transducer_stateless7/zipformer.py b/egs/librispeech/ASR/pruned_transducer_stateless7/zipformer.py
index 5b75b8d35..cbde2a2e4 100644
--- a/egs/librispeech/ASR/pruned_transducer_stateless7/zipformer.py
+++ b/egs/librispeech/ASR/pruned_transducer_stateless7/zipformer.py
@@ -281,7 +281,7 @@ class Zipformer(EncoderInterface):
 
         lengths = (x_lens - 7) >> 1
         assert x.size(0) == lengths.max().item(), (x.shape, lengths, lengths.max())
-        mask = make_pad_mask(lengths)
+        mask = make_pad_mask(lengths, x.size(0))
 
         outputs = []
         feature_masks = self.get_feature_masks(x)