refactor(amp): migrate from torch.cuda.amp to torch.amp

README.md

@@ -397,7 +397,7 @@ Few pointers to get you started:
 - [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/pytorch/ignite/blob/master/examples/notebooks/FastaiLRFinder_MNIST.ipynb) [Basic example of LR finder on
   MNIST](https://github.com/pytorch/ignite/blob/master/examples/notebooks/FastaiLRFinder_MNIST.ipynb)
 - [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/pytorch/ignite/blob/master/examples/notebooks/Cifar100_bench_amp.ipynb) [Benchmark mixed precision training on Cifar100:
-  torch.cuda.amp vs nvidia/apex](https://github.com/pytorch/ignite/blob/master/examples/notebooks/Cifar100_bench_amp.ipynb)
+  torch.amp vs nvidia/apex](https://github.com/pytorch/ignite/blob/master/examples/notebooks/Cifar100_bench_amp.ipynb)
 - [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/pytorch/ignite/blob/master/examples/notebooks/MNIST_on_TPU.ipynb) [MNIST training on a single
   TPU](https://github.com/pytorch/ignite/blob/master/examples/notebooks/MNIST_on_TPU.ipynb)
 - [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/drive/1E9zJrptnLJ_PKhmaP5Vhb6DTVRvyrKHx) [CIFAR10 Training on multiple TPUs](https://github.com/pytorch/ignite/tree/master/examples/cifar10)

docs/source/conf.py

@@ -348,7 +348,7 @@ def run(self):
     ("py:class", "torch.optim.optimizer.Optimizer"),
     ("py:class", "torch.utils.data.dataset.Dataset"),
     ("py:class", "torch.utils.data.sampler.BatchSampler"),
-    ("py:class", "torch.cuda.amp.grad_scaler.GradScaler"),
+    ("py:class", "torch.amp.grad_scaler.GradScaler"),
     ("py:class", "torch.optim.lr_scheduler._LRScheduler"),
     ("py:class", "torch.optim.lr_scheduler.LRScheduler"),
     ("py:class", "torch.utils.data.dataloader.DataLoader"),

examples/cifar10/main.py

@@ -8,7 +8,7 @@
 import torch.optim as optim
 import utils
 from torch.amp import autocast
-from torch.cuda.amp import GradScaler
+from torch.amp import GradScaler
 
 import ignite
 import ignite.distributed as idist

examples/cifar100_amp_benchmark/benchmark_torch_cuda_amp.py

@@ -1,7 +1,7 @@
 import fire
 import torch
 from torch.amp import autocast
-from torch.cuda.amp import GradScaler
+from torch.amp import GradScaler
 from torch.nn import CrossEntropyLoss
 from torch.optim import SGD
 from torchvision.models import wide_resnet50_2

examples/cifar10_qat/main.py

@@ -7,7 +7,7 @@
 import torch.optim as optim
 import utils
 from torch.amp import autocast
-from torch.cuda.amp import GradScaler
+from torch.amp import GradScaler
 
 import ignite
 import ignite.distributed as idist

examples/notebooks/Cifar100_bench_amp.ipynb

@@ -8,15 +8,15 @@
    "source": [
     "# Benchmark mixed precision training on Cifar100\n",
     "\n",
-    "In this notebook we will benchmark 1) native PyTorch mixed precision module [`torch.cuda.amp`](https://pytorch.org/docs/master/amp.html) and 2) NVidia/Apex package.\n",
+    "In this notebook we will benchmark 1) native PyTorch mixed precision module [`torch.amp`](https://pytorch.org/docs/master/amp.html) and 2) NVidia/Apex package.\n",
     "\n",
     "We will train Wide-ResNet model on Cifar100 dataset using Turing enabled GPU and compare training times.\n",
     "\n",
     "**TL;DR**\n",
     "\n",
     "The ranking is the following:\n",
     "- 1st place: Nvidia/Apex \"O2\"\n",
-    "- 2nd place: `torch.cuda.amp`: autocast and scaler\n",
+    "- 2nd place: `torch.amp`: autocast and scaler\n",
     "- 3rd place: Nvidia/Apex \"O1\"\n",
     "- 4th place: fp32\n",
     "\n",
@@ -31,7 +31,7 @@
    "source": [
     "## Installations and setup\n",
     "\n",
-    "1) Recently added [`torch.cuda.amp`](https://pytorch.org/docs/master/notes/amp_examples.html#working-with-multiple-models-losses-and-optimizers) module to perform automatic mixed precision training instead of using Nvidia/Apex package is available in PyTorch >=1.6.0.\n",
+    "1) Recently added [`torch.amp`](https://pytorch.org/docs/master/notes/amp_examples.html#working-with-multiple-models-losses-and-optimizers) module to perform automatic mixed precision training instead of using Nvidia/Apex package is available in PyTorch >=1.6.0.\n",
     "\n",
     "In this example we only need `pynvml` and `fire` packages, assuming that `torch` and `ignite` are already installed. We can install it using pip:"
    ]
@@ -154,7 +154,7 @@
     "id": "n2p-EMwGfDHs"
    },
    "source": [
-    "## Training with `torch.cuda.amp`"
+    "## Training with `torch.amp`"
    ]
   },
   {