From f6f7038145f3365ad23a0805d0cd896f0c07b3de Mon Sep 17 00:00:00 2001
From: "Yu, Guangye" <guangye.yu@intel.com>
Date: Wed, 18 Dec 2024 23:09:47 +0000
Subject: [PATCH] [4/N] Refine beginner tutorial by accelerator api

---
 .../examples_autograd/polynomial_autograd.py         |  6 +++++-
 beginner_source/fgsm_tutorial.py                     | 12 ++++--------
 beginner_source/transfer_learning_tutorial.py        |  6 +++++-
 3 files changed, 14 insertions(+), 10 deletions(-)

diff --git a/beginner_source/examples_autograd/polynomial_autograd.py b/beginner_source/examples_autograd/polynomial_autograd.py
index 9c992d2ca4..5e2fdc4101 100755
--- a/beginner_source/examples_autograd/polynomial_autograd.py
+++ b/beginner_source/examples_autograd/polynomial_autograd.py
@@ -17,8 +17,12 @@
 import torch
 import math
 
+# We want to be able to train our model on an `accelerator <https://pytorch.org/docs/stable/torch.html#accelerators>`__
+# such as CUDA, MPS, MTIA, or XPU. If the current accelerator is available, we will use it. Otherwise, we use the CPU.
+
 dtype = torch.float
-device = "cuda" if torch.cuda.is_available() else "cpu"
+device = torch.accelerator.current_accelerator().type if torch.accelerator.is_available() else "cpu"
+print(f"Using {device} device")
 torch.set_default_device(device)
 
 # Create Tensors to hold input and outputs.
diff --git a/beginner_source/fgsm_tutorial.py b/beginner_source/fgsm_tutorial.py
index 9bdf52d84b..a991fe8562 100644
--- a/beginner_source/fgsm_tutorial.py
+++ b/beginner_source/fgsm_tutorial.py
@@ -125,14 +125,9 @@
 #    `pytorch/examples/mnist <https://github.com/pytorch/examples/tree/master/mnist>`__.
 #    For simplicity, download the pretrained model `here <https://drive.google.com/file/d/1HJV2nUHJqclXQ8flKvcWmjZ-OU5DGatl/view?usp=drive_link>`__.
 #
-# -  ``use_cuda`` - boolean flag to use CUDA if desired and available.
-#    Note, a GPU with CUDA is not critical for this tutorial as a CPU will
-#    not take much time.
-#
 
 epsilons = [0, .05, .1, .15, .2, .25, .3]
 pretrained_model = "data/lenet_mnist_model.pth"
-use_cuda=True
 # Set random seed for reproducibility
 torch.manual_seed(42)
 
@@ -184,9 +179,10 @@ def forward(self, x):
             ])),
         batch_size=1, shuffle=True)
 
-# Define what device we are using
-print("CUDA Available: ",torch.cuda.is_available())
-device = torch.device("cuda" if use_cuda and torch.cuda.is_available() else "cpu")
+# We want to be able to train our model on an `accelerator <https://pytorch.org/docs/stable/torch.html#accelerators>`__
+# such as CUDA, MPS, MTIA, or XPU. If the current accelerator is available, we will use it. Otherwise, we use the CPU.
+device = torch.accelerator.current_accelerator().type if torch.accelerator.is_available() else "cpu"
+print(f"Using {device} device")
 
 # Initialize the network
 model = Net().to(device)
diff --git a/beginner_source/transfer_learning_tutorial.py b/beginner_source/transfer_learning_tutorial.py
index de7a178bd7..8a344d3d88 100644
--- a/beginner_source/transfer_learning_tutorial.py
+++ b/beginner_source/transfer_learning_tutorial.py
@@ -98,7 +98,11 @@
 dataset_sizes = {x: len(image_datasets[x]) for x in ['train', 'val']}
 class_names = image_datasets['train'].classes
 
-device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+# We want to be able to train our model on an `accelerator <https://pytorch.org/docs/stable/torch.html#accelerators>`__
+# such as CUDA, MPS, MTIA, or XPU. If the current accelerator is available, we will use it. Otherwise, we use the CPU.
+
+device = torch.accelerator.current_accelerator().type if torch.accelerator.is_available() else "cpu"
+print(f"Using {device} device")
 
 ######################################################################
 # Visualize a few images