fix triton onnx export (#1730)

k2-fsa · Aug 23, 2024 · 3b434fe · 3b434fe
1 parent 3fc06cc
commit 3b434fe
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/egs/librispeech/ASR/zipformer/export-onnx.py b/egs/librispeech/ASR/zipformer/export-onnx.py
@@ -218,7 +218,7 @@ def forward(
             - encoder_out_lens, A 1-D tensor of shape (N,)
         """
         x, x_lens = self.encoder_embed(x, x_lens)
-        src_key_padding_mask = make_pad_mask(x_lens)
+        src_key_padding_mask = make_pad_mask(x_lens, x.shape[1])
         x = x.permute(1, 0, 2)
         encoder_out, encoder_out_lens = self.encoder(x, x_lens, src_key_padding_mask)
         encoder_out = encoder_out.permute(1, 0, 2)