Fixing torch.ctc err (#1485)

* fixing torch.ctc err * Move targets & lengths to CPU
k2-fsa · Feb 2, 2024 · b9e6327 · b9e6327
1 parent b07d547
commit b9e6327
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/egs/librispeech/ASR/zipformer/model.py b/egs/librispeech/ASR/zipformer/model.py
@@ -164,9 +164,9 @@ def forward_ctc(
 
         ctc_loss = torch.nn.functional.ctc_loss(
             log_probs=ctc_output.permute(1, 0, 2),  # (T, N, C)
-            targets=targets,
-            input_lengths=encoder_out_lens,
-            target_lengths=target_lengths,
+            targets=targets.cpu(),
+            input_lengths=encoder_out_lens.cpu(),
+            target_lengths=target_lengths.cpu(),
             reduction="sum",
         )
         return ctc_loss