Fix screen extent computation return type

desy-ml · Dec 11, 2023 · ff9c409 · ff9c409
1 parent 02b0469
commit ff9c409
Showing 1 changed file with 8 additions and 6 deletions.
diff --git a/cheetah/accelerator.py b/cheetah/accelerator.py
@@ -1328,12 +1328,14 @@ def effective_pixel_size(self) -> torch.Tensor:
         return self.pixel_size * self.binning
 
     @property
-    def extent(self) -> tuple[float, float, float, float]:
-        return (
-            -self.resolution[0] * self.pixel_size[0] / 2,
-            self.resolution[0] * self.pixel_size[0] / 2,
-            -self.resolution[1] * self.pixel_size[1] / 2,
-            self.resolution[1] * self.pixel_size[1] / 2,
+    def extent(self) -> torch.Tensor:
+        return torch.stack(
+            [
+                -self.resolution[0] * self.pixel_size[0] / 2,
+                self.resolution[0] * self.pixel_size[0] / 2,
+                -self.resolution[1] * self.pixel_size[1] / 2,
+                self.resolution[1] * self.pixel_size[1] / 2,
+            ]
         )
 
     @property