fix classification model cuda move (#1039)

RyanBGG · Sep 2, 2022 · 61a53df · 61a53df
1 parent c492c9d
commit 61a53df
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/doctr/models/classification/predictor/pytorch.py b/doctr/models/classification/predictor/pytorch.py
@@ -44,7 +44,8 @@ def forward(
             raise ValueError("incorrect input shape: all pages are expected to be multi-channel 2D images.")
 
         processed_batches = self.pre_processor(crops)
-        predicted_batches = [self.model(batch) for batch in processed_batches]
+        _device = next(self.model.parameters()).device
+        predicted_batches = [self.model(batch).to(device=_device) for batch in processed_batches]
 
         # Postprocess predictions
         predicted_batches = [out_batch.argmax(dim=1).cpu().detach().numpy() for out_batch in predicted_batches]