I'm running inference_demo.sh and encountering an issue. It seems that the ProteiMPNN weight file isn't being loaded correctly. How can I resolve this：
[2026-04-21 17:06:50,034][src.utils.inference.infer_runner][INFO] - Finish loading checkpoint.
[2026-04-21 17:06:50,055][src.utils.inference.infer_runner][INFO] - Loading ProteinMPNN inverse folding module from ./ckpt/v_48_020.pt
Error executing job with overrides: ['exp=train_odesign_base_prot_flex', 'data_root_dir=./data', 'ckpt_root_dir=./ckpt', 'exp.infer_model_name=odesign_base_prot_flex', 'exp.design_modality=protein', 'exp.input_json_path=./examples/protein_design/prot_binding_prot/odesign_input.json', 'exp.exp_name=protein_binding_protein_design', 'exp.seeds=[42]', 'exp.model.sample_diffusion.N_sample=5', 'exp.use_msa=false', 'exp.num_workers=4', 'exp.model.inference_noise_schedulers.coordinate.partial_diffusion.enable=false', 'exp.model.inference_noise_schedulers.coordinate.partial_diffusion.snr=0.1']
Traceback (most recent call last):
File "/home/xcc/software/ODesign/ODesign/./scripts/inference.py", line 64, in main
infer_runner = InferRunner(
File "/home/xcc/software/ODesign/ODesign/src/utils/inference/infer_runner.py", line 50, in init
self.load_invfold_module()
File "/home/xcc/software/ODesign/ODesign/src/utils/inference/infer_runner.py", line 87, in load_invfold_module
self.model.invfold_module = build_default_wrapper(
File "/home/xcc/software/ODesign/ODesign/src/model/modules/invfold/proteinmpnn/wrapper.py", line 324, in build_default_wrapper
return MPNNFamilyWrapper(config=cfg, device=device)
File "/home/xcc/software/ODesign/ODesign/src/model/modules/invfold/proteinmpnn/wrapper.py", line 55, in init
self._load_all_models()
File "/home/xcc/software/ODesign/ODesign/src/model/modules/invfold/proteinmpnn/wrapper.py", line 58, in _load_all_models
self.models["protein_mpnn"], self.model_meta["protein_mpnn"] = self._load_model(
File "/home/xcc/software/ODesign/ODesign/src/model/modules/invfold/proteinmpnn/wrapper.py", line 87, in _load_model
model.load_state_dict(checkpoint["model_state_dict"])
File "/home/xcc/anaconda3/envs/odesign/lib/python3.10/site-packages/torch/nn/modules/module.py", line 2189, in load_state_dict
raise RuntimeError('Error(s) in loading state_dict for {}:\n\t{}'.format(
RuntimeError: Error(s) in loading state_dict for ProteinMPNN:
Unexpected key(s) in state_dict: "W_v.weight", "W_v.bias", "features.node_embedding.weight", "features.norm_nodes.weight", "features.norm_nodes.bias".
size mismatch for features.edge_embedding.weight: copying a param with shape torch.Size([128, 167]) from checkpoint, the shape in current model is torch.Size([128, 416]).
Set the environment variable HYDRA_FULL_ERROR=1 for a complete stack trace.
🎉 FINISHED: ODesign inference completed successfully!
I'm running inference_demo.sh and encountering an issue. It seems that the ProteiMPNN weight file isn't being loaded correctly. How can I resolve this：
[2026-04-21 17:06:50,034][src.utils.inference.infer_runner][INFO] - Finish loading checkpoint.
[2026-04-21 17:06:50,055][src.utils.inference.infer_runner][INFO] - Loading ProteinMPNN inverse folding module from ./ckpt/v_48_020.pt
Error executing job with overrides: ['exp=train_odesign_base_prot_flex', 'data_root_dir=./data', 'ckpt_root_dir=./ckpt', 'exp.infer_model_name=odesign_base_prot_flex', 'exp.design_modality=protein', 'exp.input_json_path=./examples/protein_design/prot_binding_prot/odesign_input.json', 'exp.exp_name=protein_binding_protein_design', 'exp.seeds=[42]', 'exp.model.sample_diffusion.N_sample=5', 'exp.use_msa=false', 'exp.num_workers=4', 'exp.model.inference_noise_schedulers.coordinate.partial_diffusion.enable=false', 'exp.model.inference_noise_schedulers.coordinate.partial_diffusion.snr=0.1']
Traceback (most recent call last):
File "/home/xcc/software/ODesign/ODesign/./scripts/inference.py", line 64, in main
infer_runner = InferRunner(
File "/home/xcc/software/ODesign/ODesign/src/utils/inference/infer_runner.py", line 50, in init
self.load_invfold_module()
File "/home/xcc/software/ODesign/ODesign/src/utils/inference/infer_runner.py", line 87, in load_invfold_module
self.model.invfold_module = build_default_wrapper(
File "/home/xcc/software/ODesign/ODesign/src/model/modules/invfold/proteinmpnn/wrapper.py", line 324, in build_default_wrapper
return MPNNFamilyWrapper(config=cfg, device=device)
File "/home/xcc/software/ODesign/ODesign/src/model/modules/invfold/proteinmpnn/wrapper.py", line 55, in init
self._load_all_models()
File "/home/xcc/software/ODesign/ODesign/src/model/modules/invfold/proteinmpnn/wrapper.py", line 58, in _load_all_models
self.models["protein_mpnn"], self.model_meta["protein_mpnn"] = self._load_model(
File "/home/xcc/software/ODesign/ODesign/src/model/modules/invfold/proteinmpnn/wrapper.py", line 87, in _load_model
model.load_state_dict(checkpoint["model_state_dict"])
File "/home/xcc/anaconda3/envs/odesign/lib/python3.10/site-packages/torch/nn/modules/module.py", line 2189, in load_state_dict
raise RuntimeError('Error(s) in loading state_dict for {}:\n\t{}'.format(
RuntimeError: Error(s) in loading state_dict for ProteinMPNN:
Unexpected key(s) in state_dict: "W_v.weight", "W_v.bias", "features.node_embedding.weight", "features.norm_nodes.weight", "features.norm_nodes.bias".
size mismatch for features.edge_embedding.weight: copying a param with shape torch.Size([128, 167]) from checkpoint, the shape in current model is torch.Size([128, 416]).
Set the environment variable HYDRA_FULL_ERROR=1 for a complete stack trace.
🎉 FINISHED: ODesign inference completed successfully!