diff --git a/invokeai/app/services/invocation_stats/invocation_stats_base.py b/invokeai/app/services/invocation_stats/invocation_stats_base.py
index b28220e74c..3266d985fe 100644
--- a/invokeai/app/services/invocation_stats/invocation_stats_base.py
+++ b/invokeai/app/services/invocation_stats/invocation_stats_base.py
@@ -30,7 +30,7 @@ writes to the system log is stored in InvocationServices.performance_statistics.
 
 from abc import ABC, abstractmethod
 from pathlib import Path
-from typing import Iterator
+from typing import ContextManager
 
 from invokeai.app.invocations.baseinvocation import BaseInvocation
 from invokeai.app.services.invocation_stats.invocation_stats_common import InvocationStatsSummary
@@ -50,7 +50,7 @@ class InvocationStatsServiceBase(ABC):
         self,
         invocation: BaseInvocation,
         graph_execution_state_id: str,
-    ) -> Iterator[None]:
+    ) -> ContextManager[None]:
         """
         Return a context object that will capture the statistics on the execution
         of invocaation. Use with: to place around the part of the code that executes the invocation.
diff --git a/invokeai/app/services/invocation_stats/invocation_stats_default.py b/invokeai/app/services/invocation_stats/invocation_stats_default.py
index 06a5b675c3..5a41f1f5d6 100644
--- a/invokeai/app/services/invocation_stats/invocation_stats_default.py
+++ b/invokeai/app/services/invocation_stats/invocation_stats_default.py
@@ -2,7 +2,7 @@ import json
 import time
 from contextlib import contextmanager
 from pathlib import Path
-from typing import Iterator
+from typing import Generator
 
 import psutil
 import torch
@@ -41,7 +41,7 @@ class InvocationStatsService(InvocationStatsServiceBase):
         self._invoker = invoker
 
     @contextmanager
-    def collect_stats(self, invocation: BaseInvocation, graph_execution_state_id: str) -> Iterator[None]:
+    def collect_stats(self, invocation: BaseInvocation, graph_execution_state_id: str) -> Generator[None, None, None]:
         # This is to handle case of the model manager not being initialized, which happens
         # during some tests.
         services = self._invoker.services