Update model/utils.py

model/utils.py

@@ -562,15 +562,15 @@ def load_checkpoint(model, ckpt_path, device, use_ema = True):
 
     ckpt_type = ckpt_path.split(".")[-1]
     if ckpt_type == "safetensors":
-        from safetensors.torch import load_file
-        checkpoint = load_file(ckpt_path, device=device)
+       # from safetensors.torch import load_file
+       # checkpoint = load_file(ckpt_path, device=device)
     else:
-        checkpoint = torch.load(ckpt_path, weights_only=False, map_location=device)
+        checkpoint = torch.load(ckpt_path, weights_only=True, map_location=device)
 
     if use_ema == True:
         ema_model = EMA(model, include_online_model = False).to(device)
         if ckpt_type == "safetensors":
-            ema_model.load_state_dict(checkpoint)
+            ema_model.load_state_dict(checkpoint['ema_model_state_dict'])
         else:
             ema_model.load_state_dict(checkpoint['ema_model_state_dict'])
         ema_model.copy_params_from_ema_to_model()