Merge pull request #24 from yfyeung/yfyeung-patch-1

Update rnnt_loss.py
k2-fsa · Jul 19, 2023 · a801adc · a801adc
2 parents 2945bd7 + 878d7c8
commit a801adc
Showing 1 changed file with 8 additions and 12 deletions.
diff --git a/fast_rnnt/python/fast_rnnt/rnnt_loss.py b/fast_rnnt/python/fast_rnnt/rnnt_loss.py
@@ -167,12 +167,10 @@ def get_rnnt_logprobs(
 
     # px is the probs of the actual symbols..
     px_am = torch.gather(
-        am.unsqueeze(1).expand(B, S, T, C),
-        dim=3,
-        index=symbols.reshape(B, S, 1, 1).expand(B, S, T, 1),
-    ).squeeze(
-        -1
-    )  # [B][S][T]
+        am.transpose(1, 2),  # (B, C, T)
+        dim=1,
+        index=symbols.unsqueeze(2).expand(B, S, T),
+    ) #  (B, S, T)
 
     if rnnt_type == "regular":
         px_am = torch.cat(
@@ -1247,12 +1245,10 @@ def get_rnnt_logprobs_smoothed(
 
     # px is the probs of the actual symbols (not yet normalized)..
     px_am = torch.gather(
-        am.unsqueeze(1).expand(B, S, T, C),
-        dim=3,
-        index=symbols.reshape(B, S, 1, 1).expand(B, S, T, 1),
-    ).squeeze(
-        -1
-    )  # [B][S][T]
+        am.transpose(1, 2),  # (B, C, T)
+        dim=1,
+        index=symbols.unsqueeze(2).expand(B, S, T),
+    ) #  (B, S, T)
 
     if rnnt_type == "regular":
         px_am = torch.cat(