Update on "[doc] fix amp doc"

# Motivation The doc shown on https://pytorch.org/docs/main/amp.html#torch.amp.autocast_mode.is_autocast_available should be `torch.amp.is_autocast_available` instead of `torch.amp.autocast_mode.is_autocast_available`. This PR aims to fix it. cc mcarilli ptrblck leslie-fang-intel jgong5 [ghstack-poisoned]
pytorch · May 11, 2024 · 56af51b · 56af51b
2 parents 3d99d26 + 35260e0
commit 56af51b
Show file tree

Hide file tree

Showing 3 changed files with 3 additions and 23 deletions.
diff --git a/docs/source/amp.rst b/docs/source/amp.rst
@@ -40,14 +40,9 @@ For CUDA and CPU, APIs are also provided separately:
 
 Autocasting
 ^^^^^^^^^^^
-.. currentmodule:: torch.amp.autocast_mode
-
-.. autofunction::  is_autocast_available
-
+.. automodule:: torch.amp.autocast_mode
 .. currentmodule:: torch.amp
-
-.. autoclass:: DOC_TEST
-    .. automethod:: test
+.. autofunction::  is_autocast_available
 
 .. currentmodule:: torch
 

diff --git a/torch/amp/__init__.py b/torch/amp/__init__.py
@@ -3,6 +3,5 @@
     _exit_autocast,
     autocast,
     is_autocast_available,
-    DOC_TEST,
 )
 from .grad_scaler import GradScaler
diff --git a/torch/amp/autocast_mode.py b/torch/amp/autocast_mode.py
@@ -6,21 +6,7 @@
 import torch
 from torch.types import _dtype
 
-__all__ = ["autocast_decorator", "autocast", "is_autocast_available", "DOC_TEST"]
-
-
-class DOC_TEST:
-    r"""
-    A test class for doc test.
-    """
-    def __init__(self) -> None:
-        pass
-
-    def test(self):
-        r"""
-        A test method.
-        """
-        print("test doc~")
+__all__ = ["autocast_decorator", "autocast", "is_autocast_available"]
 
 
 def is_autocast_available(device_type: str) -> bool: