KeyError: 'batch[0]'

def atoms_collate_fn(batch):
    print("Batch received in collate_fn:", batch)
    elem = batch[0]  # This line fails if batch is a dict or tensor
    ...

I run hotpp train in the examples/carbon, error "KeyError: '0'" when batch is a dict or tensor, even batch is null. I feel confused, please help me with this problem.

Thanks!

---

Traceback (most recent call last):
  File "/root/miniconda3/envs/quantum/bin/hotpp", line 33, in <module>
    sys.exit(load_entry_point('miaonet', 'console_scripts', 'hotpp')())
  File "/home/StarWars/hotpp/hotpp/entrypoints/main.py", line 236, in main
    f(**dict_args)
  File "/home/StarWars/hotpp/hotpp/entrypoints/train.py", line 515, in main
    trainer.fit(lit_model, datamodule=dataset)
  File "/root/miniconda3/envs/quantum/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 544, in fit
    call._call_and_handle_interrupt(
  File "/root/miniconda3/envs/quantum/lib/python3.8/site-packages/pytorch_lightning/trainer/call.py", line 43, in _call_and_handle_interrupt
    return trainer.strategy.launcher.launch(trainer_fn, *args, trainer=trainer, **kwargs)
  File "/root/miniconda3/envs/quantum/lib/python3.8/site-packages/pytorch_lightning/strategies/launchers/subprocess_script.py", line 105, in launch
    return function(*args, **kwargs)
  File "/root/miniconda3/envs/quantum/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 580, in _fit_impl
    self._run(model, ckpt_path=ckpt_path)
  File "/root/miniconda3/envs/quantum/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 987, in _run
    results = self._run_stage()
  File "/root/miniconda3/envs/quantum/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 1033, in _run_stage
    self.fit_loop.run()
  File "/root/miniconda3/envs/quantum/lib/python3.8/site-packages/pytorch_lightning/loops/fit_loop.py", line 205, in run
    self.advance()
  File "/root/miniconda3/envs/quantum/lib/python3.8/site-packages/pytorch_lightning/loops/fit_loop.py", line 363, in advance
    self.epoch_loop.run(self._data_fetcher)
  File "/root/miniconda3/envs/quantum/lib/python3.8/site-packages/pytorch_lightning/loops/training_epoch_loop.py", line 140, in run
    self.advance(data_fetcher)
  File "/root/miniconda3/envs/quantum/lib/python3.8/site-packages/pytorch_lightning/loops/training_epoch_loop.py", line 212, in advance
    batch, _, __ = next(data_fetcher)
  File "/root/miniconda3/envs/quantum/lib/python3.8/site-packages/pytorch_lightning/loops/fetchers.py", line 133, in __next__
    batch = super().__next__()
  File "/root/miniconda3/envs/quantum/lib/python3.8/site-packages/pytorch_lightning/loops/fetchers.py", line 60, in __next__
    batch = next(self.iterator)
  File "/root/miniconda3/envs/quantum/lib/python3.8/site-packages/pytorch_lightning/utilities/combined_loader.py", line 341, in __next__
    out = next(self._iterator)
  File "/root/miniconda3/envs/quantum/lib/python3.8/site-packages/pytorch_lightning/utilities/combined_loader.py", line 78, in __next__
    out[i] = next(self.iterators[i])
  File "/root/miniconda3/envs/quantum/lib/python3.8/site-packages/torch/utils/data/dataloader.py", line 628, in __next__
    data = self._next_data()
  File "/root/miniconda3/envs/quantum/lib/python3.8/site-packages/torch/utils/data/dataloader.py", line 1333, in _next_data
    return self._process_data(data)
  File "/root/miniconda3/envs/quantum/lib/python3.8/site-packages/torch/utils/data/dataloader.py", line 1359, in _process_data
    data.reraise()
  File "/root/miniconda3/envs/quantum/lib/python3.8/site-packages/torch/_utils.py", line 543, in reraise
    raise exception
KeyError: Caught KeyError in DataLoader worker process 0.
Original Traceback (most recent call last):
  File "/root/miniconda3/envs/quantum/lib/python3.8/site-packages/torch/utils/data/_utils/worker.py", line 302, in _worker_loop
    data = fetcher.fetch(index)
  File "/root/miniconda3/envs/quantum/lib/python3.8/site-packages/torch/utils/data/_utils/fetch.py", line 61, in fetch
    return self.collate_fn(data)
  File "/home/StarWars/hotpp/hotpp/data/base.py", line 126, in atoms_collate_fn
    elem = batch[0]
KeyError: 0