Update script.py

script.py

@@ -35,7 +35,9 @@ class PytorchWorker:
             model.heads.head = torch.nn.Linear(model.heads.head.in_features,
                                                number_of_categories)
 
-            model.load_state_dict(torch.load(model_path))
+            model = model.to(self.device)
+
+            model.load_state_dict(torch.load(model_path, map_location=self.device))
 
             # if not torch.cuda.is_available():
             #     model_ckpt = torch.load(model_path, map_location=torch.device("cpu"))
@@ -45,7 +47,7 @@ class PytorchWorker:
             # model_ckpt = torch.load(model_path, map_location=self.device)
             # model.load_state_dict(model_ckpt)
 
-            return model.to(self.device).eval()
+            return model.eval()
 
         self.model = _load_model(model_name, model_path)