🐛 [BUG] GMM model evaluation

[nfragapane]

Describe the bug
Trying to run nequip-evaluate --dataset-config path/to/dataset-config.yaml --model deployed_model.pth --output out.xyz --output-fields node_features_nll --output-fields-from-original-dataset forces on a deployed gmm model, but keep getting RuntimeError: Extra field node_features_nll isn't registered as node/edge/graph. Have specified node_fields: - node_features_nll in my dataset-config.yaml, so not sure why it isn't being read in correctly. Any help would be very appreciated! Full traceback:

/u/vld/quee4073/miniconda3/envs/nequip-gmm/lib/python3.8/site-packages/nequip/__init__.py:20: UserWarning: !! PyTorch version 1.13.1 found. Upstream issues in PyTorch versions 1.13.* and 2.* have been seen to cause unusual performance degredations on some CUDA systems that become worse over time; see https://github.com/mir-group/nequip/discussions/311. The best tested PyTorch version to use with CUDA devices is 1.11; while using other versions if you observe this problem, an unexpected lack of this problem, or other strange behavior, please post in the linked GitHub issue.
  warnings.warn(
Using device: cuda
Please note that _all_ machine learning models running on CUDA hardware are generally somewhat nondeterministic and that this can manifest in small, generally unimportant variation in the final test errors.
Loading model... 
    loaded deployed model
Loading dataset...
Loaded dataset specified in dataset_config.yaml.
Using all frames from the specified test dataset, yielding a test set size of 379 frames.
Starting...
  0%|                                                                                                                                                               | 0/379 [00:01<?, ?it/s]
Traceback (most recent call last):
  File "/u/vld/quee4073/miniconda3/envs/nequip-gmm/bin/nequip-evaluate", line 8, in <module>
    sys.exit(main())
  File "/u/vld/quee4073/miniconda3/envs/nequip-gmm/lib/python3.8/site-packages/nequip/scripts/evaluate.py", line 425, in main
    AtomicData.from_AtomicDataDict(output_out)
  File "/u/vld/quee4073/miniconda3/envs/nequip-gmm/lib/python3.8/site-packages/nequip/data/AtomicData.py", line 661, in to_ase
    raise RuntimeError(
RuntimeError: Extra field `node_features_nll` isn't registered as node/edge/graph

Environment (please complete the following information):

• OS: Linux
• Python 3.8.19
• python environment (commands are given for python interpreter):
  • nequip version 0.6.1
  • e3nn version 0.5.1
  • pytorch version 1.13.1