From 775a9ba1c233c286ea0f1991e9e1ef7a4ff37d30 Mon Sep 17 00:00:00 2001 From: "Adam J. Stewart" Date: Wed, 14 Aug 2024 17:16:40 +0200 Subject: [PATCH] Check torch version --- kornia/utils/helpers.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/kornia/utils/helpers.py b/kornia/utils/helpers.py index 42f5f7f286..4805afdae1 100644 --- a/kornia/utils/helpers.py +++ b/kornia/utils/helpers.py @@ -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_enabled("cpu") + 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()