Stepper errors when loading checkpoint

[cnagda]

Hi,

I'm trying to use the ACE-climSST checkpoint on hugging face (also tried the one on Zenodo), but am encountering this error:

Traceback (most recent call last):
  File "<frozen runpy>", line 198, in _run_module_as_main
  File "<frozen runpy>", line 88, in _run_code
  File "/u/cnagda2/.conda/envs/ace/lib/python3.11/site-packages/fme/ace/inference/__main__.py", line 15, in <module>
    main(args.yaml_config, segments=args.segments, override_dotlist=args.override)
  File "/u/cnagda2/.conda/envs/ace/lib/python3.11/site-packages/fme/ace/inference/inference.py", line 239, in main
    return run_inference_from_config(config)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/u/cnagda2/.conda/envs/ace/lib/python3.11/site-packages/fme/ace/inference/inference.py", line 263, in run_inference_from_config
    stepper_config = config.load_stepper_config()
                     ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/u/cnagda2/.conda/envs/ace/lib/python3.11/site-packages/fme/ace/inference/inference.py", line 204, in load_stepper_config
    return load_stepper_config(self.checkpoint_path, self.stepper_override)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/u/cnagda2/.conda/envs/ace/lib/python3.11/site-packages/fme/ace/stepper/single_module.py", line 1595, in load_stepper_config
    stepper = load_stepper(checkpoint_path, override_config)
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/u/cnagda2/.conda/envs/ace/lib/python3.11/site-packages/fme/ace/stepper/single_module.py", line 1619, in load_stepper
    stepper = Stepper.from_state(checkpoint["stepper"])
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/u/cnagda2/.conda/envs/ace/lib/python3.11/site-packages/fme/ace/stepper/single_module.py", line 1514, in from_state
    loss_normalizer = StandardNormalizer.from_state(
                      ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/u/cnagda2/.conda/envs/ace/lib/python3.11/site-packages/fme/ace/normalizer.py", line 168, in from_state
    k: torch.tensor(v, dtype=torch.float) for k, v in state["means"].items()
        ~~~~~^^^^^^^^^
TypeError: 'NoneType' object is not subscriptable
Return code: 1

This is my inference config

experiment_dir: inference_output
n_forward_steps: 24
forward_steps_in_memory: 24
checkpoint_path: ace_ckpt.tar
logging:
  log_to_screen: true
  log_to_wandb: false
  log_to_file: true
  project: ace
initial_condition:
  path: climSST/ic_2021.zarr
  engine: zarr
  start_indices:
    n_initial_conditions: 1
    first: 0
    interval: 1
forcing_loader:
  dataset:
    data_path: climSST
    file_pattern: forcing_2021.zarr
    engine: zarr
    n_repeats: 1
  num_data_workers: 0
data_writer:
  save_prediction_files: true
  save_monthly_files: false

I used pip install fme to install

[oliverwm1]

Thanks for report @cnagda. You can resolve this by pinning the fme version to before the last release, i.e. pip install fme<2025.7.0. Another option is to use a more recent model (e.g. ACE2-ERA5) which will work with the latest release. We'll also look into restoring compatibility with old checkpoints.

[oliverwm1]

Should be resolved by #46, let us know if not.