Update src/pipeline.py

src/pipeline.py

@@ -81,12 +81,14 @@ def load_pipeline() -> Pipeline:
     if encoder_path is not None:
         encoder_state_dict = torch.load(encoder_path, map_location="cpu", weights_only=True)
         filtered_state_dict = {k.strip('encoder.'): v for k, v in encoder_state_dict.items() if k.strip('encoder.') in vae.encoder.state_dict() and v.size() == vae.encoder.state_dict()[k.strip('encoder.')].size()}
+        print(f" num of keys in filtered: {len(filtered_state_dict)} and in decoder: {len(vae.encoder.state_dict())}")
         vae.encoder.load_state_dict(filtered_state_dict, strict=False)
         vae.encoder.to(dtype=dtype)
             
     if decoder_path is not None:
         decoder_state_dict = torch.load(decoder_path, map_location="cpu", weights_only=True)
         filtered_state_dict = {k.strip('decoder.'): v for k, v in decoder_state_dict.items() if k.strip('decoder.') in vae.decoder.state_dict() and v.size() == vae.decoder.state_dict()[k.strip('decoder.')].size()}
+        print(f" num of keys in filtered: {len(filtered_state_dict)} and in decoder: {len(vae.decoder.state_dict())}")
         vae.decoder.load_state_dict(filtered_state_dict, strict=False)
         vae.decoder.to(dtype=dtype)