Update src/pipeline.py

src/pipeline.py

@@ -594,8 +594,8 @@ def load_pipeline() -> Pipeline:
     ).to(memory_format=torch.channels_last)
     
     vae = AutoencoderTiny.from_pretrained("silentdriver/7815792fb4", revision="bdb7d88ebe5a1c6b02a3c0c78651dd57a403fdf5", torch_dtype=dtype)
-    vae.encoder.load_state_dict(torch.load("encoder.pth"))
-    vae.decoder.load_state_dict(torch.load("decoder.pth"))
+    vae.encoder.load_state_dict(torch.load("encoder.pth"), strict=False)
+    vae.decoder.load_state_dict(torch.load("decoder.pth"), strict=False)
     
     path = os.path.join(HF_HUB_CACHE, "models--silentdriver--7d92df966a/snapshots/add1b8d9a84c728c1209448c4a695759240bad3c")
     generator = torch.Generator(device=device)