desy-ml / cheetah

Fast and differentiable particle accelerator optics simulation for reinforcement learning and optimisation applications.
https://cheetah-accelerator.readthedocs.io
GNU General Public License v3.0
32 stars 13 forks source link

Dipole tracking throws error when using batched input #169

Closed cr-xu closed 3 months ago

cr-xu commented 3 months ago

The batched version upadte missed an if-condition in dipole traking, which causes error when batched inputs are used. https://github.com/desy-ml/cheetah/blob/df0844b1820da8a05c62f80d6c9bb98a8e0887ed/cheetah/accelerator/dipole.py#L123