Skip to content

Commit 80989d2

Browse files
guangyeysvekars
andauthored
[WIP][4/N] Refine beginner tutorial by accelerator api (#3190)
Co-authored-by: Svetlana Karslioglu <[email protected]>
1 parent 15f7cb3 commit 80989d2

File tree

3 files changed

+14
-10
lines changed

3 files changed

+14
-10
lines changed

beginner_source/examples_autograd/polynomial_autograd.py

+5-1
Original file line numberDiff line numberDiff line change
@@ -17,8 +17,12 @@
1717
import torch
1818
import math
1919

20+
# We want to be able to train our model on an `accelerator <https://pytorch.org/docs/stable/torch.html#accelerators>`__
21+
# such as CUDA, MPS, MTIA, or XPU. If the current accelerator is available, we will use it. Otherwise, we use the CPU.
22+
2023
dtype = torch.float
21-
device = "cuda" if torch.cuda.is_available() else "cpu"
24+
device = torch.accelerator.current_accelerator().type if torch.accelerator.is_available() else "cpu"
25+
print(f"Using {device} device")
2226
torch.set_default_device(device)
2327

2428
# Create Tensors to hold input and outputs.

beginner_source/fgsm_tutorial.py

+4-8
Original file line numberDiff line numberDiff line change
@@ -125,14 +125,9 @@
125125
# `pytorch/examples/mnist <https://github.com/pytorch/examples/tree/master/mnist>`__.
126126
# For simplicity, download the pretrained model `here <https://drive.google.com/file/d/1HJV2nUHJqclXQ8flKvcWmjZ-OU5DGatl/view?usp=drive_link>`__.
127127
#
128-
# - ``use_cuda`` - boolean flag to use CUDA if desired and available.
129-
# Note, a GPU with CUDA is not critical for this tutorial as a CPU will
130-
# not take much time.
131-
#
132128

133129
epsilons = [0, .05, .1, .15, .2, .25, .3]
134130
pretrained_model = "data/lenet_mnist_model.pth"
135-
use_cuda=True
136131
# Set random seed for reproducibility
137132
torch.manual_seed(42)
138133

@@ -184,9 +179,10 @@ def forward(self, x):
184179
])),
185180
batch_size=1, shuffle=True)
186181

187-
# Define what device we are using
188-
print("CUDA Available: ",torch.cuda.is_available())
189-
device = torch.device("cuda" if use_cuda and torch.cuda.is_available() else "cpu")
182+
# We want to be able to train our model on an `accelerator <https://pytorch.org/docs/stable/torch.html#accelerators>`__
183+
# such as CUDA, MPS, MTIA, or XPU. If the current accelerator is available, we will use it. Otherwise, we use the CPU.
184+
device = torch.accelerator.current_accelerator().type if torch.accelerator.is_available() else "cpu"
185+
print(f"Using {device} device")
190186

191187
# Initialize the network
192188
model = Net().to(device)

beginner_source/transfer_learning_tutorial.py

+5-1
Original file line numberDiff line numberDiff line change
@@ -98,7 +98,11 @@
9898
dataset_sizes = {x: len(image_datasets[x]) for x in ['train', 'val']}
9999
class_names = image_datasets['train'].classes
100100

101-
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
101+
# We want to be able to train our model on an `accelerator <https://pytorch.org/docs/stable/torch.html#accelerators>`__
102+
# such as CUDA, MPS, MTIA, or XPU. If the current accelerator is available, we will use it. Otherwise, we use the CPU.
103+
104+
device = torch.accelerator.current_accelerator().type if torch.accelerator.is_available() else "cpu"
105+
print(f"Using {device} device")
102106

103107
######################################################################
104108
# Visualize a few images

0 commit comments

Comments
 (0)