🐛 [BUG] Smaller dataset when using initialize_from_state

[wwz171]

Bug Description

When I try to train a pre-trained model on a smaller dataset, I get the following error:

Traceback (most recent call last):
  File "/home/weizhouwang/.conda/envs/NEQUIP/bin/nequip-train", line 10, in <module>
    sys.exit(main())
  File "/home/weizhouwang/.conda/envs/NEQUIP/lib/python3.9/site-packages/nequip/scripts/train.py", line 110, in main
    trainer = fresh_start(config)
  File "/home/weizhouwang/.conda/envs/NEQUIP/lib/python3.9/site-packages/nequip/scripts/train.py", line 240, in fresh_start
    trainer.set_dataset(dataset, validation_dataset)
  File "/home/weizhouwang/.conda/envs/NEQUIP/lib/python3.9/site-packages/nequip/train/trainer.py", line 1263, in set_dataset
    self.dataset_train = dataset.index_select(self.train_idcs)
  File "/home/weizhouwang/.conda/envs/NEQUIP/lib/python3.9/site-packages/nequip/utils/torch_geometric/dataset.py", line 220, in index_select
    return self.index_select(idx.flatten().tolist())
  File "/home/weizhouwang/.conda/envs/NEQUIP/lib/python3.9/site-packages/nequip/utils/torch_geometric/dataset.py", line 234, in index_select
    indices = [indices[i] for i in idx]
  File "/home/weizhouwang/.conda/envs/NEQUIP/lib/python3.9/site-packages/nequip/utils/torch_geometric/dataset.py", line 234, in <listcomp>
    indices = [indices[i] for i in idx]
IndexError: range object index out of range

The model is trained on a bigger dataset, but I want to fine-tune it on a smaller dataset. And it goes wrong.

I guess the error is caused by config.yaml file of key include_file_as_baseline_config: there's a train_idcs and val_idcs in it which indicate the indices of the training and validation data. When I train the model on a smaller dataset, the indices are out of range. After I delete these two lines, the error is gone.