diff --git a/src/python/gem5/simulate/simulator.py b/src/python/gem5/simulate/simulator.py
index b32e54aaad..1d0d3ecc66 100644
--- a/src/python/gem5/simulate/simulator.py
+++ b/src/python/gem5/simulate/simulator.py
@@ -87,7 +87,7 @@ class Simulator:
         board: AbstractBoard,
         full_system: Optional[bool] = None,
         on_exit_event: Optional[
-            Dict[Union[str, ExitEvent], Generator[Optional[bool], None, None]]
+            Dict[ExitEvent, Generator[Optional[bool], None, None]]
         ] = None,
         expected_execution_order: Optional[List[ExitEvent]] = None,
         checkpoint_path: Optional[Path] = None,