Skip to content

Commit 4b3f897

Browse files
authored
Merge branch 'main' into guangyey/accelerator4
2 parents 253f2e2 + b9b1656 commit 4b3f897

File tree

8 files changed

+847
-219
lines changed

8 files changed

+847
-219
lines changed

.jenkins/validate_tutorials_built.py

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -25,7 +25,6 @@
2525
"intermediate_source/mnist_train_nas", # used by ax_multiobjective_nas_tutorial.py
2626
"intermediate_source/fx_conv_bn_fuser",
2727
"intermediate_source/_torch_export_nightly_tutorial", # does not work on release
28-
"intermediate_source/transformer_building_blocks", # does not work on release
2928
"advanced_source/super_resolution_with_onnxruntime",
3029
"advanced_source/usb_semisup_learn", # fails with CUDA OOM error, should try on a different worker
3130
"prototype_source/fx_graph_mode_ptq_dynamic",

beginner_source/basics/buildmodel_tutorial.py

Lines changed: 4 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -32,17 +32,10 @@
3232
#############################################
3333
# Get Device for Training
3434
# -----------------------
35-
# We want to be able to train our model on a hardware accelerator like the GPU or MPS,
36-
# if available. Let's check to see if `torch.cuda <https://pytorch.org/docs/stable/notes/cuda.html>`_
37-
# or `torch.backends.mps <https://pytorch.org/docs/stable/notes/mps.html>`_ are available, otherwise we use the CPU.
38-
39-
device = (
40-
"cuda"
41-
if torch.cuda.is_available()
42-
else "mps"
43-
if torch.backends.mps.is_available()
44-
else "cpu"
45-
)
35+
# We want to be able to train our model on an `accelerator <https://pytorch.org/docs/stable/torch.html#accelerators>`__
36+
# such as CUDA, MPS, MTIA, or XPU. If the current accelerator is available, we will use it. Otherwise, we use the CPU.
37+
38+
device = torch.accelerator.current_accelerator().type if torch.accelerator.is_available() else "cpu"
4639
print(f"Using {device} device")
4740

4841
##############################################

en-wordlist.txt

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -159,6 +159,7 @@ Minifier
159159
MobileNet
160160
ModelABC
161161
MPS
162+
MTIA
162163
Mypy
163164
NAS
164165
NCCL

0 commit comments

Comments
 (0)