Apply suggestions from code review

pytorch · Dec 9, 2024 · 13d95db · 13d95db
1 parent 36dda87
commit 13d95db
Showing 1 changed file with 8 additions and 2 deletions.
diff --git a/ignite/engine/engine.py b/ignite/engine/engine.py
@@ -998,7 +998,10 @@ def _internal_run_as_gen(self) -> Generator[Any, None, State]:
                         # update time wrt handlers
                         self.state.times[Events.EPOCH_COMPLETED.name] = epoch_time_taken
 
-                    self.should_terminate_single_epoch = False
+                    if self.should_terminate_single_epoch:
+                        # We skip raising _EngineTerminateSingleEpochException exception on Events.EPOCH_COMPLETED
+                        # as epoch is already completed and nothing to terminate 
+                        self.should_terminate_single_epoch = False
                     yield from self._maybe_terminate_or_interrupt()
 
                     hours, mins, secs = _to_hours_mins_secs(epoch_time_taken)
@@ -1192,7 +1195,10 @@ def _internal_run_legacy(self) -> State:
                         # update time wrt handlers
                         self.state.times[Events.EPOCH_COMPLETED.name] = epoch_time_taken
 
-                    self.should_terminate_single_epoch = False
+                    if self.should_terminate_single_epoch:
+                        # We skip raising _EngineTerminateSingleEpochException exception on Events.EPOCH_COMPLETED
+                        # as epoch is already completed and nothing to terminate 
+                        self.should_terminate_single_epoch = False
                     self._maybe_terminate_legacy()
 
                     hours, mins, secs = _to_hours_mins_secs(epoch_time_taken)