Update rnnt_loss.py

4bff8cc2 · Yifan Yang · GitHub · 445aca29 · 4bff8cc2
Unverified Commit 4bff8cc2 authored May 03, 2023 by Yifan Yang Committed by GitHub May 03, 2023
Hide whitespace changes
Inline Side-by-side

Showing with 2 additions and 2 deletions

fast_rnnt/python/fast_rnnt/rnnt_loss.py fast_rnnt/python/fast_rnnt/rnnt_loss.py +2 -2

No files found.
--- a/fast_rnnt/python/fast_rnnt/rnnt_loss.py
+++ b/fast_rnnt/python/fast_rnnt/rnnt_loss.py
@@ -773,9 +773,9 @@ def do_rnnt_pruning(
        lm,
        dim=1,
        index=ranges.reshape(B, T * s_range, 1).expand(
-            (B, T * s_range, decoder_dim)
+            (B, T * s_range, C)
        ),
-    ).reshape(B, T, s_range, decoder_dim)
+    ).reshape(B, T, s_range, C)
    return am_pruning, lm_pruning