fix: weights_only=False for SelVA checkpoints (PyTorch 2.6 compat)

PyTorch 2.6 changed the default to weights_only=True. SelVA checkpoints
contain non-tensor types (numpy scalars etc.) that fail strict unpickling.
All weights come from trusted sources (jnwnlee/selva HF repo).

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
2026-04-04 16:38:31 +02:00
parent 40388ba6de
commit 614a2e02aa
4 changed files with 5 additions and 5 deletions
+1 -1
View File
@@ -19,7 +19,7 @@ class AutoEncoderModule(nn.Module):
need_vae_encoder: bool = True):
super().__init__()
self.vae: VAE = get_my_vae(mode).eval()
vae_state_dict = torch.load(vae_ckpt_path, weights_only=True, map_location='cpu')
vae_state_dict = torch.load(vae_ckpt_path, weights_only=False, map_location='cpu')
self.vae.load_state_dict(vae_state_dict)
self.vae.remove_weight_norm()