diff --git a/beginner_source/blitz/cifar10_tutorial.py b/beginner_source/blitz/cifar10_tutorial.py index 8f19f5964c6..008f29315fe 100644 --- a/beginner_source/blitz/cifar10_tutorial.py +++ b/beginner_source/blitz/cifar10_tutorial.py @@ -77,13 +77,15 @@ trainset = torchvision.datasets.CIFAR10(root='./data', train=True, download=True, transform=transform) +# num_workers=0 avoids multiprocessing spawn issues when running this file as +# ``python cifar10_tutorial.py`` on macOS/Windows (see note above). trainloader = torch.utils.data.DataLoader(trainset, batch_size=batch_size, - shuffle=True, num_workers=2) + shuffle=True, num_workers=0) testset = torchvision.datasets.CIFAR10(root='./data', train=False, download=True, transform=transform) testloader = torch.utils.data.DataLoader(testset, batch_size=batch_size, - shuffle=False, num_workers=2) + shuffle=False, num_workers=0) classes = ('plane', 'car', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck') @@ -292,27 +294,30 @@ def forward(self, x): ######################################################################## # Okay, so what next? # -# How do we run these neural networks on the GPU? +# How do we run these neural networks on a GPU or other accelerator? # # Training on GPU # ---------------- -# Just like how you transfer a Tensor onto the GPU, you transfer the neural -# net onto the GPU. +# Just like how you transfer a Tensor onto a device, you transfer the neural +# net onto that device. # -# Let's first define our device as the first visible cuda device if we have -# CUDA available: - -device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu') - -# Assuming that we are on a CUDA machine, this should print a CUDA device: - +# Let's first select a device. This picks the fastest available accelerator, +# or falls back to CPU (same pattern as ``quickstart_tutorial.py``). +device = ( + torch.accelerator.current_accelerator().type + if torch.accelerator.is_available() + else "cpu" +) + +# ``device`` is a string such as ``"cuda"``, ``"mps"``, or ``"cpu"``, which you +# can pass to ``Tensor.to(...)`` and ``nn.Module.to(...)``. print(device) ######################################################################## -# The rest of this section assumes that ``device`` is a CUDA device. +# The snippets below show how to move the model and batch data to ``device``. # # Then these methods will recursively go over all modules and convert their -# parameters and buffers to CUDA tensors: +# parameters and buffers to tensors on ``device``: # # .. code:: python # @@ -320,7 +325,7 @@ def forward(self, x): # # # Remember that you will have to send the inputs and targets at every step -# to the GPU too: +# to ``device`` too: # # .. code:: python #