Update vae_2_2.py

fab43a07 · gushiqiao · GitHub · 7516ad2a · fab43a07
Commit fab43a07 authored Aug 12, 2025 by gushiqiao Committed by GitHub Aug 12, 2025
Hide whitespace changes
Inline Side-by-side

Showing with 1 addition and 1 deletion

lightx2v/models/video_encoders/hf/wan/vae_2_2.py lightx2v/models/video_encoders/hf/wan/vae_2_2.py +1 -1

No files found.
--- a/lightx2v/models/video_encoders/hf/wan/vae_2_2.py
+++ b/lightx2v/models/video_encoders/hf/wan/vae_2_2.py
@@ -828,7 +828,7 @@ def _video_vae(pretrained_path=None, z_dim=16, dim=160, device="cpu", cpu_offloa
    # load checkpoint
    logging.info(f"loading {pretrained_path}")
    weights_dict = load_weights(pretrained_path, cpu_offload=cpu_offload)
-    model.load_state_dict(weights_dict)
+    model.load_state_dict(weights_dict, assign=True)

    return model