Update skyreels_v2_infer/modules/vae.py

skyreels_v2_infer/modules/vae.py

@@ -563,7 +563,7 @@ def _video_vae(pretrained_path=None, z_dim=None, device="cpu", **kwargs):
 
     # load checkpoint
     logging.info(f"loading {pretrained_path}")
-    model.load_state_dict(torch.load(pretrained_path, map_location=device), assign=True)
+    model.load_state_dict(torch.load(pretrained_path, map_location=device, weigths_only=True), assign=True)
 
     return model