Fix for rnnt_loss.py (#1177)

* Update rnnt_loss.py * Update rnnt_loss.py * Fix for style check * Fix for style check * Update rnnt_loss.py
k2-fsa · Apr 26, 2023 · a23383c · a23383c
1 parent 8c9044e
commit a23383c
Showing 1 changed file with 8 additions and 12 deletions.
diff --git a/k2/python/k2/rnnt_loss.py b/k2/python/k2/rnnt_loss.py
@@ -169,12 +169,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(
@@ -1399,12 +1397,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(