Update base.py

PaddlePaddle · Aug 6, 2024 · 6d13b9f · 6d13b9f
1 parent 8189aea
commit 6d13b9f
Showing 1 changed file with 3 additions and 4 deletions.
diff --git a/python/paddle/base/dygraph/base.py b/python/paddle/base/dygraph/base.py
@@ -21,8 +21,6 @@
     TYPE_CHECKING,
     Any,
     Callable,
-    ContextManager,
-    Sequence,
     TypeVar,
     overload,
 )
@@ -41,8 +39,9 @@
 
 if TYPE_CHECKING:
     from collections import OrderedDict
+    from collections.abc import Generator, Sequence
+    from contextlib import AbstractAsyncContextManager
     from types import TracebackType
-    from typing import Generator
 
     from typing_extensions import Self
 
@@ -299,7 +298,7 @@ def _switch_tracer_mode_guard_(
 
 
 @overload
-def no_grad(func: None = ...) -> ContextManager:
+def no_grad(func: None = ...) -> AbstractAsyncContextManager:
     ...