[no ci] map device location to cpu/cuda

src/ocrtoolkit/integrations/doctr.py

@@ -21,11 +21,12 @@ def __init__(self, model, path, device, **kwargs):
         from doctr.models.detection.predictor import DetectionPredictor
 
         super().__init__(model, path, device)
-        self.doctr_base_predictor = _OCRPredictor()
         kwargs["mean"] = kwargs.get("mean", model.cfg["mean"])
         kwargs["std"] = kwargs.get("std", model.cfg["std"])
         kwargs["batch_size"] = kwargs.get("batch_size", 2)
         input_shape = model.cfg["input_shape"][1:]
+
+        self.doctr_base_predictor = _OCRPredictor()
         self.predictor = DetectionPredictor(PreProcessor(input_shape, **kwargs), model)
 
     def _predict(self, images: List[np.ndarray], **kwargs) -> List[DetectionResults]:

src/ocrtoolkit/integrations/gcv.py

@@ -22,6 +22,9 @@ class GCVModel(DetectionModel):
     def __init__(self, client, path):
         super().__init__(model=client, path=path)
 
+    def _map_location(self):
+        pass
+
     def _predict(self, images: List[np.ndarray], **kwargs) -> List[DetectionResults]:
         l_results = []
         for image in images:

src/ocrtoolkit/wrappers/model.py

@@ -19,6 +19,10 @@ def __init__(self, model, path=None, device="cpu", **kwargs):
         self.model = model
         self.path = path
         self.device = device
+        self._map_location()
+
+    def _map_location(self):
+        self.model = self.model.to(self.device)
         logger.info(f"Loaded model from {self.path}, to {self.device}")
 
     def preprocess(self, images: List[np.ndarray], **kwargs) -> List[np.ndarray]: