torch/autograd/grad_mode.pyi

from typing import Any, Callable, TypeVar

# Used for annotating the decorator usage of 'no_grad' and 'enable_grad'.
# See https://mypy.readthedocs.io/en/latest/generics.html#declaring-decorators
FuncType = Callable[..., Any]
T = TypeVar('T', bound=FuncType)

class no_grad:
    def __enter__(self) -> None: ...
    def __exit__(self, *args: Any) -> bool: ...
    def __call__(self, func: T) -> T: ...

class enable_grad:
    def __enter__(self) -> None: ...
    def __exit__(self, *args: Any) -> bool: ...
    def __call__(self, func: T) -> T: ...

class set_grad_enabled:
    def __init__(self, mode: bool) -> None: ...
    def __enter__(self) -> None: ...
    def __exit__(self, *args: Any) -> bool: ...