Standardise "Introduction to PyTorch" to use GPU, MPS or CPU (#2276)

alvaropp · web-flow · commit 600e14a78286 · 2023-04-03T15:34:52.000-07:00
The [quickstart](https://pytorch.org/tutorials/beginner/basics/quickstart_tutorial.html) tutorial does try to use GPU/MPS if available. Later, the [Build the Neural Network](https://pytorch.org/tutorials/beginner/basics/buildmodel_tutorial.html#) tutorial only looks for GPU. This got me confused for while as I went back to find the code to select MPS and couldn't find it, and I reckon it's better to have these two tutorials be consistent.
diff --git a/beginner_source/basics/buildmodel_tutorial.py b/beginner_source/basics/buildmodel_tutorial.py
@@ -32,12 +32,17 @@
 #############################################
 # Get Device for Training
 # -----------------------
-# We want to be able to train our model on a hardware accelerator like the GPU,
-# if it is available. Let's check to see if
-# `torch.cuda <https://pytorch.org/docs/stable/notes/cuda.html>`_ is available, else we
-# continue to use the CPU.
-
-device = "cuda" if torch.cuda.is_available() else "cpu"
+# We want to be able to train our model on a hardware accelerator like the GPU or MPS,
+# if available. Let's check to see if `torch.cuda <https://pytorch.org/docs/stable/notes/cuda.html>`_
+# or `torch.backends.mps <https://pytorch.org/docs/stable/notes/mps.html>`_ are available, otherwise we use the CPU.
+
+device = (
+    "cuda"
+    if torch.cuda.is_available()
+    else "mps"
+    if torch.backends.mps.is_available()
+    else "cpu"
+)
 print(f"Using {device} device")
 
 ##############################################
diff --git a/beginner_source/basics/quickstart_tutorial.py b/beginner_source/basics/quickstart_tutorial.py
@@ -84,10 +84,16 @@
 # To define a neural network in PyTorch, we create a class that inherits
 # from `nn.Module <https://pytorch.org/docs/stable/generated/torch.nn.Module.html>`_. We define the layers of the network
 # in the ``__init__`` function and specify how data will pass through the network in the ``forward`` function. To accelerate
-# operations in the neural network, we move it to the GPU if available.
-
-# Get cpu or gpu device for training.
-device = "cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu"
+# operations in the neural network, we move it to the GPU or MPS if available.
+
+# Get cpu, gpu or mps device for training.
+device = (
+    "cuda"
+    if torch.cuda.is_available()
+    else "mps"
+    if torch.backends.mps.is_available()
+    else "cpu"
+)
 print(f"Using {device} device")
 
 # Define model
