Update frame_reducer.py

kobenaxie · Dec 7, 2023 · 10f65e1 · 10f65e1
1 parent 2be6fc3
commit 10f65e1
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/egs/librispeech/ASR/pruned_transducer_stateless7_ctc_bs/frame_reducer.py b/egs/librispeech/ASR/pruned_transducer_stateless7_ctc_bs/frame_reducer.py
@@ -79,9 +79,11 @@ def forward(
             limit_lens = x_lens - y_lens
             max_limit_len = limit_lens.max().int()
             ctc_prob = ctc_output.exp()
+            ctc_prob = ctc_prob.masked_fill(
+                padding_mask.unsqueeze(-1), 0.0,
+            )
             fake_limit_indexes = torch.topk(
-                ctc_prob[:, :, blank_id].masked_fill(padding_mask, 0.0), 
-                k=max_limit_len,
+                ctc_prob[:, :, blank_id], max_limit_len
             ).indices
             T_arange = (
                 torch.arange(max_limit_len)