diff --git a/kornia/utils/helpers.py b/kornia/utils/helpers.py index 3738ee21c6..4805afdae1 100644 --- a/kornia/utils/helpers.py +++ b/kornia/utils/helpers.py @@ -305,7 +305,7 @@ def is_autocast_enabled(both: bool = True) -> bool: Returns: Return a Bool, will always return False for a torch without support, otherwise will be: if both is True - `torch.is_autocast_enabled() or torch.is_autocast_cpu_enabled()`. If both is False will return just + `torch.is_autocast_enabled() or torch.is_autocast_enabled('cpu')`. If both is False will return just `torch.is_autocast_enabled()`. """ if TYPE_CHECKING: @@ -316,7 +316,10 @@ def is_autocast_enabled(both: bool = True) -> bool: return False if both: - return torch.is_autocast_enabled() or torch.is_autocast_cpu_enabled() + if torch_version_ge(2, 4): + return torch.is_autocast_enabled() or torch.is_autocast_enabled("cpu") + else: + return torch.is_autocast_enabled() or torch.is_autocast_cpu_enabled() return torch.is_autocast_enabled()