cpu

CozmoxJC · Jul 20, 2022 · 9c98a72 · 9c98a72
1 parent 585cdba
commit 9c98a72
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/utils.py b/utils.py
@@ -26,7 +26,7 @@ def load_ckpt(model, ckpt_path, model_name='model', prefixes_to_ignore=[]):
 
 
 def slim_ckpt(ckpt_path):
-    ckpt = torch.load(ckpt_path)
+    ckpt = torch.load(ckpt_path, map_location='cpu')
     # pop unused parameters
     keys_to_pop = []
     for k in ckpt['state_dict']: