From 3b434fe83c40eaf3c4739c26d11aa3a3b8af8ddc Mon Sep 17 00:00:00 2001 From: Yuekai Zhang Date: Fri, 23 Aug 2024 09:33:46 +0800 Subject: [PATCH] fix triton onnx export (#1730) --- egs/librispeech/ASR/zipformer/export-onnx.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/egs/librispeech/ASR/zipformer/export-onnx.py b/egs/librispeech/ASR/zipformer/export-onnx.py index ed8a0ef0fd..ca3cbf0d59 100755 --- 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)