Merge pull request #26 from yfyeung/patch-1

Update rnnt_loss.py

Merge pull request #26 from yfyeung/patch-1
Update rnnt_loss.py
2945bd7d · Daniel Povey · GitHub · 2c2dc4b9 · 4bff8cc2 · 2945bd7d
Unverified Commit 2945bd7d authored Jul 19, 2023 by Daniel Povey Committed by GitHub Jul 19, 2023
Hide whitespace changes
Inline Side-by-side

Showing with 6 additions and 4 deletions

fast_rnnt/python/fast_rnnt/rnnt_loss.py fast_rnnt/python/fast_rnnt/rnnt_loss.py +6 -4

No files found.
--- a/fast_rnnt/python/fast_rnnt/rnnt_loss.py
+++ b/fast_rnnt/python/fast_rnnt/rnnt_loss.py
@@ -770,10 +770,12 @@ def do_rnnt_pruning(
    # (B, T, s_range, C)
    lm_pruning = torch.gather(
-        lm.unsqueeze(1).expand((B, T, S + 1, C)),
+        lm,
-        dim=2,
+        dim=1,
-        index=ranges.reshape((B, T, s_range, 1)).expand((B, T, s_range, C)),
+        index=ranges.reshape(B, T * s_range, 1).expand(
-    )
+            (B, T * s_range, C)
+        ),
+    ).reshape(B, T, s_range, C)
    return am_pruning, lm_pruning