Add type annotations

desy-ml · Dec 4, 2024 · cfb362e · cfb362e
1 parent adb59bb
commit cfb362e
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/cheetah/particles/particle_beam.py b/cheetah/particles/particle_beam.py
@@ -1,4 +1,4 @@
-from typing import Optional
+from typing import Optional, Union
 
 import torch
 from scipy import constants
@@ -1074,7 +1074,7 @@ def momenta(self) -> torch.Tensor:
         """Momenta of the individual particles."""
         return torch.sqrt(self.energies**2 - electron_mass_eV**2)
 
-    def __getitem__(self, item):
+    def __getitem__(self, item: Union[int, slice, torch.Tensor]) -> "ParticleBeam":
         vector_shape = torch.broadcast_shapes(
             self.particles.shape[:-2],
             self.energy.shape,