Check torch version

CodSpeedHQ · Aug 14, 2024 · 775a9ba · 775a9ba
1 parent 1e467be
commit 775a9ba
Showing 1 changed file with 4 additions and 1 deletion.
diff --git 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()