pytorch · guangyey · May 7, 2024 · May 7, 2024 · May 7, 2024 · May 7, 2024
diff --git a/docs/source/amp.rst b/docs/source/amp.rst
@@ -40,8 +40,6 @@ For CUDA and CPU, APIs are also provided separately:
 
 Autocasting
 ^^^^^^^^^^^
-.. currentmodule:: torch.amp.autocast_mode
-
 .. autofunction::  is_autocast_available
 
 .. currentmodule:: torch

diff --git a/torch/amp/__init__.py b/torch/amp/__init__.py
@@ -2,6 +2,8 @@
     _enter_autocast,
     _exit_autocast,
     autocast,
-    is_autocast_available,
+    is_autocast_available as is_autocast_available,
 )
 from .grad_scaler import GradScaler
+
+__all__ = ["GradScaler", "autocast", "is_autocast_available"]