Skip to content

Commit

Permalink
[2/N] Refine beginner tutorial by accelerator api
Browse files Browse the repository at this point in the history
  • Loading branch information
guangyey committed Nov 26, 2024
1 parent 540bd0c commit 90ed2df
Show file tree
Hide file tree
Showing 2 changed files with 12 additions and 18 deletions.
14 changes: 4 additions & 10 deletions beginner_source/basics/quickstart_tutorial.py
Original file line number Diff line number Diff line change
Expand Up @@ -84,16 +84,10 @@
# 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 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"
)
# operations in the neural network, we move it to the `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")

# Define model
Expand Down
16 changes: 8 additions & 8 deletions beginner_source/basics/tensorqs_tutorial.py
Original file line number Diff line number Diff line change
Expand Up @@ -99,20 +99,20 @@
# Operations on Tensors
# ~~~~~~~~~~~~~~~~~~~~~~~
#
# Over 100 tensor operations, including arithmetic, linear algebra, matrix manipulation (transposing,
# Over 1200 tensor operations, including arithmetic, linear algebra, matrix manipulation (transposing,
# indexing, slicing), sampling and more are
# comprehensively described `here <https://pytorch.org/docs/stable/torch.html>`__.
#
# Each of these operations can be run on the GPU (at typically higher speeds than on a
# CPU). If you’re using Colab, allocate a GPU by going to Runtime > Change runtime type > GPU.
# Each of these operations can be run on the CPU and `Accelerator <https://pytorch.org/docs/stable/torch.html#accelerators>`__
# such as CUDA, MPS, MTIA, or XPU. If you’re using Colab, allocate an accelerator by going to Runtime > Change runtime type > GPU.
#
# By default, tensors are created on the CPU. We need to explicitly move tensors to the GPU using
# ``.to`` method (after checking for GPU availability). Keep in mind that copying large tensors
# By default, tensors are created on the CPU. We need to explicitly move tensors to the accelerator using
# ``.to`` method (after checking for accelerator availability). Keep in mind that copying large tensors
# across devices can be expensive in terms of time and memory!

# We move our tensor to the GPU if available
if torch.cuda.is_available():
tensor = tensor.to("cuda")
# We move our tensor to the current accelerator if available
if torch.accelerator.is_available():
tensor = tensor.to(torch.accelerator.current_accelerator())


######################################################################
Expand Down

0 comments on commit 90ed2df

Please sign in to comment.