Fix engine mypy issue

pytorch · Jul 31, 2024 · bd8e88b · bd8e88b
1 parent 1e13f2c
commit bd8e88b
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/ignite/engine/engine.py b/ignite/engine/engine.py
@@ -157,7 +157,7 @@ def __init__(self, process_function: Callable[["Engine", Any], Any]):
         _check_signature(process_function, "process_function", self, None)
 
         # generator provided by self._internal_run_as_gen
-        self._internal_run_generator: Optional[Generator] = None
+        self._internal_run_generator: Optional[Generator[Any, None, State]] = None
 
     def register_events(
         self, *event_names: Union[List[str], List[EventEnum]], event_to_attr: Optional[dict] = None
@@ -951,7 +951,7 @@ def _internal_run(self) -> State:
             self._internal_run_generator = None
             return out.value
 
-    def _internal_run_as_gen(self) -> Generator:
+    def _internal_run_as_gen(self) -> Generator[Any, None, State]:
         self.should_terminate = self.should_terminate_single_epoch = self.should_interrupt = False
         self._init_timers(self.state)
         try: