Error: KeyError: 'graph_encoder.layers.0.encoder.lin_src.weight'

Dear authors, Thank you for your great work on Fairseq. When I tried to run the below command, I got this error. Hopefully I would received some advices from you.

🐛 Bug

I got the bug KeyError: 'graph_encoder.layers.0.encoder.lin_src.weight' when trying to load checkpoint_utils.load_model_ensemble from fairseq.

To Reproduce

from fairseq import options, tasks, checkpoint_utils, progress_bar, utils
parser = options.get_generation_parser()
args = options.parse_args_and_arch(parser)
task = tasks.setup_task(args)
models, _ = checkpoint_utils.load_model_ensemble(["models/generation_models/tplgen_template_generation/model.pt"], task=task)

where the models folder it is from https://drive.google.com/drive/folders/1EpfamTiOosKvy_s9Bm_tW1NdtAzKpkY7 and I alreadly used the model so I think it works fine.

Expected behavior

models, _ = checkpoint_utils.load_model_ensemble([args.path], task=task)
  File "/content/gdrive/.shortcut-targets-by-id/1euwk1Zyt7Sa5sxm4_AVBZ9m4ECrtcIo7/A*/open-ended_question_ontology/fairseq/fairseq/checkpoint_utils.py", line 368, in load_model_ensemble
    state,
  File "/content/gdrive/.shortcut-targets-by-id/1euwk1Zyt7Sa5sxm4_AVBZ9m4ECrtcIo7/A*/open-ended_question_ontology/fairseq/fairseq/checkpoint_utils.py", line 471, in load_model_ensemble_and_task
    state["model"], strict=strict, model_cfg=cfg.model
  File "/content/gdrive/.shortcut-targets-by-id/1euwk1Zyt7Sa5sxm4_AVBZ9m4ECrtcIo7/A*/open-ended_question_ontology/fairseq/fairseq/models/fairseq_model.py", line 128, in load_state_dict
    return super().load_state_dict(new_state_dict, strict)
  File "/usr/local/lib/python3.7/dist-packages/torch/nn/modules/module.py", line 1468, in load_state_dict
    load(self)
  File "/usr/local/lib/python3.7/dist-packages/torch/nn/modules/module.py", line 1466, in load
    load(child, prefix + name + '.')
  File "/usr/local/lib/python3.7/dist-packages/torch/nn/modules/module.py", line 1466, in load
    load(child, prefix + name + '.')
  File "/usr/local/lib/python3.7/dist-packages/torch/nn/modules/module.py", line 1466, in load
    load(child, prefix + name + '.')
  [Previous line repeated 2 more times]
  File "/usr/local/lib/python3.7/dist-packages/torch/nn/modules/module.py", line 1463, in load
    state_dict, prefix, local_metadata, True, missing_keys, unexpected_keys, error_msgs)
  File "/usr/local/lib/python3.7/dist-packages/torch/nn/modules/module.py", line 1372, in _load_from_state_dict
    hook(state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs)
  File "/usr/local/lib/python3.7/dist-packages/torch_geometric/nn/dense/linear.py", line 131, in _lazy_load_hook
    weight = state_dict[prefix + 'weight']
KeyError: 'graph_encoder.layers.0.encoder.lin_src.weight'

Environment

PyG version (torch_geometric.__version__): 2.0.3
PyTorch version: (torch.__version__): 1.10.0
OS (e.g., Linux): Colab, Yes
Python version (e.g., 3.9): 3.7.12
CUDA/cuDNN version: 11.1
How you installed PyTorch and PyG (conda, pip, source): I install by pip from the instruction here: https://pytorch.org/

How do you install fairseq?

git clone https://github.com/pytorch/fairseq.git
cd fairseq
python setup.py build develop
pip install --editable ./
python setup.py build_ext --inplace

Any other relevant information (e.g., version of torch-scatter): I install torch_geometric:

!pip install torch-scatter -f https://data.pyg.org/whl/torch-1.10.0+cu111.html
!pip install torch-sparse -f https://data.pyg.org/whl/torch-1.10.0+cu111.html
!pip install torch-geometric

facebookresearch / fairseq