Update model/utils.py

model/utils.py

@@ -570,8 +570,7 @@ def load_checkpoint(model, ckpt_path, device, dtype=None, use_ema=True):
 
         checkpoint = load_file(ckpt_path)
     else:
-        checkpoint = torch.load(ckpt_path, weights_only=True)
-
+        checkpoint = torch.load(ckpt_path, weights_only=True, map_location=torch.device('cpu')) 
     if use_ema:
         if ckpt_type == "safetensors":
             checkpoint = {"ema_model_state_dict": checkpoint}