TypeError: 'AnemoiModelInterface' object is not subscriptable

[Timotheevin]

What happened?

I am trying to finetune the AIFS v1 model. When I'm running the comand :

anemoi-training train --config-name=config_finetuning.yaml

Everything is workig fine until it tries to load the checkpoint and throws this error message :

conda/envs/playground/lib/python3.9/site-packages/pytorch_lightning/utilities/migration/utils.py", line 164, in _get_version
return checkpoint["pytorch-lightning_version"]
TypeError: 'AnemoiModelInterface' object is not subscriptable

What are the steps to reproduce the bug?

Run the command :
anemoi-training train --config-name=config_finetuning.yaml
with the config_finetuning.yaml file given in this huggingface repo : https://huggingface.co/ecmwf/aifs-single-1.0
(checkpoint also available there)

Version

I am using : - python3.9 - torch==2.6.0 - anemoi-training==0.3.1 - anemoi-models==0.4.0 - anemoi-graphs==0.4.4 - pytorch-lightning==2.5.1.post0

Platform (OS and architecture)

Rocky Linux release 8.10 (Green Obsidian) (Running on an HPC)

Relevant log output

File "/home/tv330/.conda/envs/playground/lib/python3.9/site-packages/anemoi/training/train/train.py", line 150, in model
    return GraphForecaster.load_from_checkpoint(self.last_checkpoint, **kwargs)
  File "/home/tv330/.conda/envs/playground/lib/python3.9/site-packages/pytorch_lightning/utilities/model_helpers.py", line 125, in wrapper
    return self.method(cls, *args, **kwargs)
  File "/home/tv330/.conda/envs/playground/lib/python3.9/site-packages/pytorch_lightning/core/module.py", line 1581, in load_from_checkpoint
    loaded = _load_from_checkpoint(
  File "/home/tv330/.conda/envs/playground/lib/python3.9/site-packages/pytorch_lightning/core/saving.py", line 68, in _load_from_checkpoint
    checkpoint = _pl_migrate_checkpoint(
  File "/home/tv330/.conda/envs/playground/lib/python3.9/site-packages/pytorch_lightning/utilities/migration/utils.py", line 143, in _pl_migrate_checkpoint
    old_version = _get_version(checkpoint)
  File "/home/tv330/.conda/envs/playground/lib/python3.9/site-packages/pytorch_lightning/utilities/migration/utils.py", line 164, in _get_version
    return checkpoint["pytorch-lightning_version"]
TypeError: 'AnemoiModelInterface' object is not subscriptable

Accompanying data

No response

Organisation

No response