problem with running i-pi

[DavideTisi]

i am facing problem running pet with i-pi.
Both for pet and i-pi I am using the main branch.

I am getting this as error:

Traceback (most recent call last):
  File "/home/tisi/ipi-test-paper/i-pi/drivers/py/driver.py", line 237, in <module>
    run_driver(
  File "/home/tisi/ipi-test-paper/i-pi/drivers/py/driver.py", line 130, in run_driver
    pot, force, vir, extras = driver(cell, pos)
  File "/home/tisi/ipi-test-paper/i-pi/drivers/py/pes/pet.py", line 94, in __call__
    pot, force = self.pet_calc.forward(pet_structure)
  File "/home/tisi/miniconda3/envs/pet/lib/python3.9/site-packages/pet/single_struct_calculator.py", line 70, in forward
    prediction_energy, prediction_forces = self.model(
  File "/home/tisi/miniconda3/envs/pet/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1518, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/home/tisi/miniconda3/envs/pet/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1527, in _call_impl
    return forward_call(*args, **kwargs)
TypeError: forward() got an unexpected keyword argument 'augmentation'

augmetation is passed here https://github.com/lab-cosmo/pet/blob/d10e31d2056f1bface44bd2eeb5e63387c727f44/src/single_struct_calculator.py#L71