fix amp doc

ghstack-source-id: 48445cc521730517f3ea483ff04cd31c2a0388b3 Pull Request resolved: #125673
pytorch · May 7, 2024 · 15a69b6 · 15a69b6
1 parent 5dee462
commit 15a69b6
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/docs/source/amp.rst b/docs/source/amp.rst
@@ -6,6 +6,7 @@ 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.amp
 .. py:module:: torch.cpu.amp
 .. py:module:: torch.cuda.amp
 
@@ -40,7 +41,7 @@ For CUDA and CPU, APIs are also provided separately:
 
 Autocasting
 ^^^^^^^^^^^
-.. currentmodule:: torch.amp.autocast_mode
+.. currentmodule:: torch.amp
 
 .. autofunction::  is_autocast_available