diff --git a/src/python/gem5/simulate/simulator.py b/src/python/gem5/simulate/simulator.py
index b32e54a..1d0d3ec 100644
--- a/src/python/gem5/simulate/simulator.py
+++ b/src/python/gem5/simulate/simulator.py
@@ -87,7 +87,7 @@
         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,