Fix padding issues (#1303)

k2-fsa · Oct 11, 2023 · 2b3c5d7 · 2b3c5d7
1 parent 16a2748
commit 2b3c5d7
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/egs/librispeech/ASR/pruned_transducer_stateless2/conformer.py b/egs/librispeech/ASR/pruned_transducer_stateless2/conformer.py
@@ -158,7 +158,7 @@ def forward(
         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
@@ -281,7 +281,7 @@ def forward(
 
         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)