envs

utils/utils.py

@@ -229,7 +229,7 @@ def load_checkpoint(model_file, model):
 def load_model(model, model_path):
     if model_path != "":
         print(f"init model from checkpoint: {model_path}")
-        model_ckpt = torch.load(model_path, map_location="cuda")
+        model_ckpt = torch.load(model_path, map_location="cpu")
         if "global_step" in model_ckpt: print(f"global_step: {model_ckpt['global_step']}")
         state_dict = model_ckpt["state_dict"] if "state_dict" in model_ckpt else model_ckpt
         m, u = model.load_state_dict(state_dict, strict=False)