Change the syntax a bit.

AngelFP · Jul 5, 2024 · 8eea77a · 8eea77a
1 parent bf6b4e6
commit 8eea77a
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/wake_t/beamline_elements/plasma_stage.py b/wake_t/beamline_elements/plasma_stage.py
@@ -112,8 +112,9 @@ def __init__(
         if self.wakefield is not None:
             fields.append(self.wakefield)
         fields.extend(self.external_fields)
-        self.auto_dt_bunch = auto_dt_bunch
-        if self.auto_dt_bunch is None:
+        if auto_dt_bunch is not None:
+            self.auto_dt_bunch = auto_dt_bunch
+        else:
             self.auto_dt_bunch = self._get_optimized_dt
         super().__init__(
             length=length,