search

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
28 stars 13 forks source link

Inconsistent device setting for Beam classes #19

Closed cr-xu closed 1 year ago

cr-xu commented 1 year ago

The Element and ParticleBeam classes have default device to "auto" which basically sets device = "cuda" if torch.cuda.is_available() else "cpu". c.f. this line and this line

The ParameterBeam, however, does not have this setting.

It will lead to RuntimeError complaining that the lattice and beam are on different devices if one runs cheetah using ParameterBeam on a PC equipped with cuda. This is quite confusing for new users.

Suggestion: either we add to the ParameterBeam also the auto device setting (don't know if it makes sense as probably there won't be any speed gain); or perform some device checking when the simulation is called and throw a warning.

cr-xu commented 1 year ago

fixed by #32 Now throws an error if tracking with CPU Paramterbeam and the lattice elements are in cuda.