Small fix to cuda use

Trusted-AI · Jan 26, 2024 · bc3c432 · bc3c432
1 parent 55b22f6
commit bc3c432
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/art/estimators/speech_recognition/pytorch_icefall.py b/art/estimators/speech_recognition/pytorch_icefall.py
@@ -160,7 +160,7 @@ def predict(self, x: np.ndarray, batch_size: int = 1, **kwargs) -> np.ndarray:
 
  # extract features
  x, _, _ = self.transform_model_input(x=torch.tensor(wav))
- shape = torch.tensor([x.shape[1]])
+ shape = torch.tensor([x.shape[1]]).to(self.device)
 
  encoder_out, encoder_out_lens = self.transducer_model.encoder(x=x, x_lens=shape)
  hyp = greedy_search(model=self.transducer_model, encoder_out=encoder_out, id2word=self.get_id2word)