Skip to content

Commit

Permalink
support onnx export with batch size 1; also works for batch processin…
Browse files Browse the repository at this point in the history
…g, but the results might be affected by the padding
  • Loading branch information
marcoyang1998 committed Apr 7, 2024
1 parent f3e8e42 commit 01b744f
Showing 1 changed file with 1 addition and 4 deletions.
5 changes: 1 addition & 4 deletions egs/audioset/AT/zipformer/export-onnx.py
Original file line number Diff line number Diff line change
Expand Up @@ -202,10 +202,7 @@ def forward(
logits = self.classifier(encoder_out) # (N, T, num_classes)
# Note that this is slightly different from model.py for better
# support of onnx
N = logits.shape[0]
for i in range(N):
logits[i, encoder_out_lens[i] :] = 0
logits = logits.sum(dim=1) / encoder_out_lens.unsqueeze(-1)
logits = logits.mean(dim=1)
return logits


Expand Down

0 comments on commit 01b744f

Please sign in to comment.