map model to device using map_location parameter

thearyadev · Feb 29, 2024 · 6187715 · 6187715
1 parent 145fac3
commit 6187715
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/heroes/hero_comparison.py b/heroes/hero_comparison.py
@@ -145,7 +145,7 @@ def predict_hero_name(self, image: Image, model_directory: Path) -> str:
 
 
         if model_directory.name not in model_cache.keys():
-            model_cache[model_directory.name] = torch.load(model_directory / "model.pth")
+            model_cache[model_directory.name] = torch.load(model_directory / "model.pth", map_location=device)
         st_dict = model_cache[model_directory.name]