From 21eca78e47612a639eed866e5d46c168d66c0a56 Mon Sep 17 00:00:00 2001 From: "Yu, Guangye" Date: Wed, 15 May 2024 19:32:22 +0000 Subject: [PATCH] fix amp doc ghstack-source-id: ae2957d91375f4f3a10d7faebe5442b24337c21b Pull Request resolved: https://github.com/pytorch/pytorch/pull/125673 --- docs/source/amp.rst | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/docs/source/amp.rst b/docs/source/amp.rst index a9d98a0aa0989..f5647afdc5814 100644 --- a/docs/source/amp.rst +++ b/docs/source/amp.rst @@ -6,8 +6,6 @@ Automatic Mixed Precision package - torch.amp .. Both modules below are missing doc entry. Adding them here for now. .. This does not add anything to the rendered page -.. py:module:: torch.cpu.amp -.. py:module:: torch.cuda.amp .. automodule:: torch.amp .. currentmodule:: torch.amp @@ -40,9 +38,11 @@ For CUDA and CPU, APIs are also provided separately: Autocasting ^^^^^^^^^^^ -.. currentmodule:: torch.amp.autocast_mode +.. autosummary:: + :toctree: generated + :nosignatures: -.. autofunction:: is_autocast_available + is_autocast_available .. currentmodule:: torch @@ -394,6 +394,8 @@ regardless of whether autocast is enabled. .. This module needs to be documented. Adding here in the meantime .. for tracking purposes +.. py:module:: torch.cpu.amp +.. py:module:: torch.cuda.amp .. py:module:: torch.amp.autocast_mode .. py:module:: torch.cpu.amp.autocast_mode .. py:module:: torch.cuda.amp.autocast_mode