is_autocast_enabled: fix deprecation warning

CodSpeedHQ · Aug 13, 2024 · 1e467be · 1e467be
1 parent 929d1b2
commit 1e467be
Showing 1 changed file with 2 additions and 2 deletions.
diff --git 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,7 @@ def is_autocast_enabled(both: bool = True) -> bool:
         return False
 
     if both:
-        return torch.is_autocast_enabled() or torch.is_autocast_cpu_enabled()
+        return torch.is_autocast_enabled() or torch.is_autocast_enabled("cpu")
 
     return torch.is_autocast_enabled()