Add device option for ParameterBeam

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

27 stars 12 forks source link

Add device option for ParameterBeam #69

Closed cr-xu closed 10 months ago

cr-xu commented 10 months ago

Now one can also specify the device for ParameterBeam
- Added the according option in the doc string
- Fixed the Element.call device check

cr-xu commented 10 months ago

On a side note, we probably should include device error checks in tracking in the pytest. This requires then CUDA / MPS device (once #61 is addressed) in CI/CD runners. I don't know how to do it at the moment though...

jank324 commented 10 months ago

Ufff ... I wonder if especially MPS is even possible. I guess we can keep this in mind and otherwise postpone it until it comes up again. But testing would probably be good because clearly there were errors now that we didn't catch.