Fixed model load on CPU-only machines

lovaslin · Jun 17, 2024 · 5021795 · 5021795
1 parent 534cc44
commit 5021795
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/src/defect_detection/functions.py b/src/defect_detection/functions.py
@@ -51,6 +51,9 @@ def deepAE_load(path, use_only=True, loss_fn=None, opt=None, opt_param=None):
     ae = AE_cls(param, use_only, loss_fn, opt, opt_param)
 
     # Load trained parameters
-    ae.load_state_dict(torch.load(path + "AE_state.save"))
+    if torch.cuda.is_available():
+        ae.load_state_dict(torch.load(path + "AE_state.save"))
+    else:
+        ae.load_state_dict(torch.load(path + "AE_state.save", map_location='cpu'))
 
     return ae