[FIX] incorrect variable name

MannLabs · Feb 2, 2025 · 9815bc0 · 9815bc0
1 parent 1b630ba
commit 9815bc0
Showing 1 changed file with 5 additions and 3 deletions.
diff --git a/src/scportrait/pipeline/classification.py b/src/scportrait/pipeline/classification.py
@@ -1201,7 +1201,7 @@ def _load_model(self):
 
         model = ConvNextModel.from_pretrained("facebook/convnext-xlarge-224-22k")
         model.eval()
-        model.to(self.device)
+        model.to(self.inference_device)
 
         return model
 
@@ -1338,7 +1338,7 @@ def inference(self, dataloader, model_fun, partial=False):
             batch_size = self.batch_size
 
             images, label, class_id = next(data_iter)
-            images["pixel_values"] = images["pixel_values"][0].to(self.device)
+            images["pixel_values"] = images["pixel_values"][0].to(self.inference_device)
             o = model_fun(**images)
             result = o.pooler_output.cpu().detach()
 
@@ -1364,7 +1364,9 @@ def inference(self, dataloader, model_fun, partial=False):
 
             for i in tqdm(range(len(dataloader) - 1)):
                 images, label, class_id = next(data_iter)
-                images["pixel_values"] = images["pixel_values"][0].to(self.device)
+                images["pixel_values"] = images["pixel_values"][0].to(
+                    self.inference_device
+                )
 
                 o = model_fun(**images)
                 result = o.pooler_output.cpu().detach()