KeyError 'stress' when using keymapping

[mstapelberg]

Hi,

I'm using:
nequip 0.6.0 develop version
allegro 0.3.0 develop version

I'm struggling with keymapping my extxyz datasets, where I have tried the following for my config file:

dataset: ase                                                # type of data set, can be npz or ase
dataset_file_name: ./data/ase_vcrti_neb_all_ind_samp_25per_train.extxyz
ase_args:
  format: extxyz
dataset_key_mapping:
  REF_energy: total_energy
  REF_force: forces
  REF_stress: stress

validation_dataset: ase
validation_dataset_file_name: ./data/ase_vcrti_neb_all_ind_samp_25per_val.extxyz
validation_ase_args:
  format: extxyz
validation_dataset_key_mapping:
  REF_energy: total_energy
  REF_force: forces
  REF_stress: stress

I see the following in the logs:
For training:

Loaded data: Batch(atomic_numbers=[1397907, 1], batch=[1397907], cell=[22189, 3, 3], edge_cell_shift=[83772862, 3], edge_index=[2, 83772862], forces=[1397907, 3], free_energy=[22189], pbc=[22189, 3], pos=[1397907, 3], ptr=[22190], total_energy=[22189, 1])
    processed data size: ~3283.10 MB

Training optional args:

ASEDataset_param = dict(
...   optional_args = {'ase_args': {'format': 'extxyz'}, 'file_name': './data/ase_vcrti_neb_all_ind_samp_25per_val.extxyz', 'url': None, 'AtomicData_options': {'r_max': 6.0}, 'include_frames': None, 'key_mapping': {'REF_energy': 'total_energy', 'REF_force': 'forces', 'REF_stress': 'stress'}, 'include_keys': ['stress'], 'root': 'results/vcrti'},
...   positional_args = {'type_mapper': <nequip.data.transforms.TypeMapper object at 0x7fa7f65320e0>})

For validation:

Loaded data: Batch(atomic_numbers=[174699, 1], batch=[174699], cell=[2773, 3, 3], edge_cell_shift=[10463362, 3], edge_index=[2, 10463362], forces=[174699, 3], free_energy=[2773], pbc=[2773, 3], pos=[174699, 3], ptr=[2774], total_energy=[2773, 1])
    processed data size: ~410.07 MB

Validation optional args:

AtomwiseReduce_param = dict(
...   optional_args = {'out_field': 'total_energy', 'reduce': 'sum', 'avg_num_atoms': None, 'field': 'atomic_energy'},
...   positional_args = {'irreps_in': {'pos': 1x1o, 'edge_index': None, 'edge_embedding': 128x0e, 'edge_cutoff': 1x0e, 'edge_attrs': 1x0e+1x1o+1x2e, 'edge_features': 1024x0e, 'edge_energy': 1x0e, 'atomic_energy': 1x0e}})

I get the following error:

  File "/home/myless/.mambaforge/envs/allegro/lib/python3.10/site-packages/nequip/train/_loss.py", line 80, in __call__
    ref = ref[key]
KeyError: 'stress'

It seems like my key_mapping for stress is not working? Has anyone had this happen to them?

Thanks,
Myles