support gpu decoding for mask-ctc

chintu619 · May 18, 2022 · 9c83ddb · 9c83ddb
1 parent 49100e4
commit 9c83ddb
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/espnet2/asr/maskctc_model.py b/espnet2/asr/maskctc_model.py
@@ -314,7 +314,10 @@ def forward(self, enc_out: torch.Tensor) -> List[Hypothesis]:
         confident_idx = torch.nonzero(probs_hat[y_idx] >= p_thres).squeeze(-1)
         mask_num = len(mask_idx)
 
-        y_in = torch.zeros(1, len(y_idx), dtype=torch.long) + self.mask_token
+        y_in = (
+            torch.zeros(1, len(y_idx), dtype=torch.long).to(enc_out.device)
+            + self.mask_token
+        )
         y_in[0][confident_idx] = y_hat[y_idx][confident_idx]
 
         logging.info("msk:{}".format(self.ids2text(y_in[0].tolist())))