from typing import Callable
import functools
from typing import Any
import torch


def torch_no_grad(func: Callable) -> Callable:
    """Decorator that runs the decorated function in a context with disabled PyTorch
    gradient calculations.

    Disabling gradient calculations is useful for inference time, when you are sure
    that you will not call `Tensor.backward()`.

    This will reduce memory consumption for computations that would otherwise have
    `requires_grad=True`. In this mode, the result of every computation will have
    `requires_grad=False`, even when the inputs have `requires_grad=True`.

    Arguments:
        func: the function to be decorated.

    Returns:
        The decorated function, which evaluates in a context with disabled gradient
        calculations.
    """

    @functools.wraps(func)
    def wrapper(*args, **kwargs) -> Any:
        """Wrapped function to be returned by the decorator.

        Returns:
            The original function evaluation.
        """
        with torch.no_grad():
            return func(*args, **kwargs)

    return wrapper