fix: add device parameter and locate model and data to device #16

lombokai · May 21, 2024 · 8fbe6bc · 8fbe6bc
1 parent 8f254e9
commit 8fbe6bc
Showing 1 changed file with 4 additions and 3 deletions.
diff --git a/src/onepiece_classify/infer/recognition.py b/src/onepiece_classify/infer/recognition.py
@@ -12,8 +12,9 @@
 
 
 class ImageRecognition(BaseInference):
-    def __init__(self, model_path: str):
+    def __init__(self, model_path: str, device: str):
         self.model_path = Path(model_path)
+        self.device = device
         self.class_dict = {
             0: "Ace",
             1: "Akainu",
@@ -39,7 +40,7 @@ def __init__(self, model_path: str):
 
     def _build_model(self):
         # load model
-        state_dict = torch.load(self.model_path)
+        state_dict = torch.load(self.model_path, map_location=self.device)
         model_backbone = image_recog(self.nclass)
         model_backbone.load_state_dict(state_dict)
         return model_backbone
@@ -65,7 +66,7 @@ def pre_process(
         else:
             print("Image type not recognized")
 
-        return img
+        return img.to(self.device)
 
     def forward(self, image_tensor: torch.Tensor) -> torch.Tensor:
         self.model.eval()