PyTorch Fundamentals

Chamil Jay

PyTorch is one of the most widely used frameworks for building and training neural networks. At the centre of PyTorch is the tensor.

A tensor is a general-purpose numerical data structure that can represent anything from a single number to a high-dimensional collection of values. Neural-network inputs, model weights, activations, gradients, images, audio representations, and many other objects are represented using tensors.

If you are using an Apple Silicon Mac, PyTorch can also take advantage of Apple’s GPU through the Metal Performance Shaders (MPS) backend.

This tutorial develops a practical understanding of PyTorch tensors by working through the operations you will encounter repeatedly when building machine-learning models.

Contents

Important: The Python code in this tutorial is kept exactly as provided. Outputs involving random values, PyTorch versions, MPS availability, and execution time will naturally vary between machines and runs.

⋅ • ✦ • ⋅

Getting Started

We begin by importing PyTorch and Python’s built-in time module.

torch provides the tensor functionality and neural-network capabilities we will use throughout the tutorial.

time will be used later to measure how long tensor operations take on the CPU and MPS.

import time
import torch

The next cell checks the PyTorch version and whether the MPS backend is available.

print("PyTorch version:", torch.__version__)
print("MPS available:", torch.backends.mps.is_available())

The exact output depends on the PyTorch installation and computer.

For example, on an Apple Silicon Mac with MPS enabled, it will look like:

### Output
PyTorch version: 2.xx
MPS available: True

The important result is the second line. If it is True, PyTorch can use Apple’s GPU through the MPS backend.

⋅ • ✦ • ⋅

Device Setup - CPU vs MPS

A tensor does not just contain values. It also has a device associated with it.

The device determines where the tensor is stored and where operations involving that tensor are performed.

The most common devices are:

  • cpu — the CPU
  • mps — Apple’s GPU backend
  • cuda — NVIDIA GPUs, on supported systems

For an Apple Silicon Mac, mps is the device we are interested in.

Rather than assuming MPS is available, we can select it conditionally:

device = "mps" if torch.backends.mps.is_available() else "cpu"
print("Selected device:", device)
sample = torch.ones(1, device=device)
sample

### Output
Selected device: mps
tensor([1.], device='mps:0')

The important concept here is that the tensor is created directly on the selected device.

For example:

CPU:
tensor stored in CPU memory

MPS:
tensor stored in memory accessible to Apple's GPU

This distinction becomes important when we start working with neural networks. The model and its input tensors generally need to be on compatible devices.

⋅ • ✦ • ⋅

Simple CPU vs GPU Benchmark

One of the reasons to use a GPU is that many machine-learning operations can be performed in parallel.

A CPU is designed to be very flexible and efficient for a wide range of workloads. A GPU, on the other hand, is designed to perform large numbers of similar numerical operations simultaneously.

This makes GPUs particularly useful for the large matrix and tensor operations that occur in deep learning.

The following function repeatedly doubles a large tensor and measures how long the operation takes.

def benchmark_tensor_add(target_device: str, size: int = 2000, steps: int = 200):
    a = torch.ones(size, size, device=target_device)
    if target_device == "mps":
        torch.mps.synchronize()
    start = time.time()
    for _ in range(steps):
        a += a
    if target_device == "mps":
        torch.mps.synchronize()
    return time.time() - start

Use of torch.mps.synchronize() is particularly important as GPU operations can be asynchronous. Python can submit an operation to the GPU and continue executing before the GPU has finished.

Without synchronisation, our timer could stop before the GPU has completed the actual computation.

Therefore:

torch.mps.synchronize()

waits for outstanding MPS operations to finish.

Now we run the benchmark on the CPU and, if available, on MPS.

cpu_time = benchmark_tensor_add("cpu")
print(f"CPU time: {cpu_time:.4f}s")

if torch.backends.mps.is_available():
    mps_time = benchmark_tensor_add("mps")
    print(f"MPS time: {mps_time:.4f}s")
else:
    print("MPS is unavailable on this machine.")

### Output
CPU time: 0.0285s
MPS time: 0.0128s

The actual numbers depend on:

  • Apple Silicon generation
  • PyTorch version
  • tensor size
  • system load
  • background applications
  • thermal conditions
  • MPS implementation

GPU acceleration is most useful when the workload contains enough parallel computation to overcome the overhead of using the GPU.

For very small tensors, the CPU can sometimes be faster.

⋅ • ✦ • ⋅

Basic Tensor Intuition

A tensor can be thought of as a generalisation of familiar mathematical structure matrix. In the essence, a tensor is a multi-dimensional array of numbers.

A useful progression is:

scalar       → 0 dimensions
vector       → 1 dimension
matrix       → 2 dimensions
higher tensor → 3+ dimensions

For example, the following is a two-dimensional tensor:

test_tensor = torch.tensor([[1, 2], [3, 4], [5, 6]])
test_tensor

### Output
tensor([[1, 2],
        [3, 4],
        [5, 6]])

This tensor has three rows and two columns.

print("ndim:", test_tensor.ndim)
print("shape:", test_tensor.shape)
print("first column:", test_tensor[:, 0])

### Output
ndim: 2
shape: torch.Size([3, 2])
first column: tensor([1, 3, 5])

⋅ • ✦ • ⋅

Creating Common Tensors

PyTorch provides many convenient functions for creating tensors.

Some of the most commonly used are:

# creates a random tensor with shape (1, 2, 3)
rand_tensor = torch.rand(1, 2, 3). 

# creates a tensor filled with zeroes with shape (5, 3)
zeros = torch.zeros(5, 3) 

# creates a tensor filled with ones with shape (5, 4)
ones = torch.ones(5, 4) 

# creates a tensor with values from 1 to 9 (10 is excluded)
range_tensor = torch.arange(1, 10, 1) 

# creates a tensor filled with zeroes with the same shape as `ones`
like_tensor = torch.zeros_like(ones) 

rand_tensor.shape, zeros.shape, ones.shape, range_tensor.shape, like_tensor.shape

### Output
(torch.Size([1, 2, 3]),
 torch.Size([5, 3]),
 torch.Size([5, 4]),
 torch.Size([9]),
 torch.Size([5, 4]))

⋅ • ✦ • ⋅

Dtype and Device

Two practical tensor properties that we will talk about today:

  • dtype controls how values are represented
  • device determines where the tensor is stored

These concepts become increasingly important when training neural networks.

float32_tensor_cpu = torch.tensor([[2, 2], [3, 4], [5, 6]], dtype=torch.float32, device="cpu")

if torch.backends.mps.is_available():
    float32_tensor_mps = torch.tensor([[2, 2], [3, 4], [5, 6]], dtype=torch.float32, device="mps")
else:
    float32_tensor_mps = None

float16_tensor_cpu = float32_tensor_cpu.to(torch.float16)

print("CPU float32 dtype/device:", float32_tensor_cpu.dtype, float32_tensor_cpu.device)
print("CPU float16 dtype/device:", float16_tensor_cpu.dtype, float16_tensor_cpu.device)
print("MPS tensor:", float32_tensor_mps)

On an Apple Silicon Mac with MPS available, the output will have this form:

### Output
CPU float32 dtype/device: torch.float32 cpu
CPU float16 dtype/device: torch.float16 cpu
MPS tensor: tensor([[2., 2.],
        [3., 4.],
        [5., 6.]], device='mps:0')

dtype means data type.

For example:

torch.float32

represents 32-bit floating-point numbers.

The choice of data type affects:

  • memory consumption
  • numerical precision
  • computational performance
  • compatibility with operations and hardware

Deep-learning models often use floating-point numbers because neural networks perform large numbers of numerical calculations involving weights, activations, and gradients. If you have heard about model quantization, it is a technique that reduces the precision of the model’s weights and activations to save memory and speed up inference.

device means device. which tells PyTorch where the tensor is located.

For example:

cpu

means the tensor is on the CPU.

While:

mps

means the tensor is using Apple’s GPU backend.

A common source of PyTorch errors is accidentally attempting to perform an operation between tensors on different devices.

⋅ • ✦ • ⋅

Tensor Arithmetic and Matrix Multiplication

Tensor arithmetic is fundamental to neural networks.

It is important to distinguish between element-wise operations and matrix multiplication.

First, consider basic arithmetic:

base = torch.tensor([[1, 2], [3, 4]])
print("add 10:\n", base + 10)
print("mul 10:\n", base * 10)
print("sub 10:\n", base - 10)

### Output
add 10:
 tensor([[11, 12],
        [13, 14]])
mul 10:
 tensor([[10, 20],
        [30, 40]])
sub 10:
 tensor([[-9, -8],
        [-7, -6]])

These operations are element-wise.

Now we compare element-wise multiplication with matrix multiplication. Between two tensors we can do element-wise multiplication using * or matrix multiplication using @, torch.matmul(), or torch.mm(), but of course the dimensions must be compatible.

x = torch.tensor([[1, 2], [3, 4]])
y = torch.tensor([[10, 20], [30, 40]])

print("Element-wise x * y:\n", x * y)
print("Matrix product matmul(x, y):\n", torch.matmul(x, y))
print("Matrix product x @ y:\n", x @ y)
print("Matrix product mm(x, y):\n", torch.mm(x, y))
print("Transpose x.T:\n", x.T)

### Output
Element-wise x * y:
 tensor([[ 10,  40],
        [ 90, 160]])
Matrix product matmul(x, y):
 tensor([[ 70, 100],
        [150, 220]])
Matrix product x @ y:
 tensor([[ 70, 100],
        [150, 220]])
Matrix product mm(x, y):
 tensor([[ 70, 100],
        [150, 220]])
Transpose x.T:
 tensor([[1, 3],
        [2, 4]])

Element-wise multiplication multiplies corresponding elements of the two tensors, while matrix multiplication follows the rules of linear algebra.

x * y ## Element-wise multiplication
x @ y ## Matrix multiplication
torch.mm(x, y) ## Matrix multiplication

⋅ • ✦ • ⋅

Reduction Operations: max, min, argmax, mean

A reduction operation summarizes values along one or more dimensions.

Common examples include: maximum, minimum, average, sum, index of maximum, index of minimum

x = torch.rand(2, 3)
print('x:', x)
print('x.max:', x.max())
print('x.min:', x.min())
print('x.argmax:', x.argmax())
print('x.argmin:', x.argmin())
print('x.mean:', x.mean())

### Output - will vary each time
x: tensor([[0.4320, 0.9194, 0.4892],
        [0.4380, 0.1121, 0.0427]])
x.max: tensor(0.9194)
x.min: tensor(0.0427)
x.argmax: tensor(1)
x.argmin: tensor(5)
x.mean: tensor(0.4056)

⋅ • ✦ • ⋅

Reshape, View, Stack, Cat, Squeeze, Unsqueeze, Permute

Tensor shape manipulation is one of the most important practical skills in PyTorch. Deep-learning models often require data to have a particular shape.

For example, an image model may expect:

batch × channels × height × width

while another framework or data source may provide:

batch × height × width × channels

Understanding how to transform these shapes is essential.

x = torch.arange(0, 16)
reshaped = x.reshape(2, 2, 2, 2)
reshaped_auto = x.reshape(2, 2, -1)
viewed = x.view(2, 2, -1)

print("x shape:", x.shape)
print("reshape(2,2,2,2) shape:", reshaped.shape)
print("reshape(2,2,-1) shape:", reshaped_auto.shape)
print("view(2,2,-1) shape:", viewed.shape)

### Output
x shape: torch.Size([16])
reshape(2,2,2,2) shape: torch.Size([2, 2, 2, 2])
reshape(2,2,-1) shape: torch.Size([2, 2, 4])
view(2,2,-1) shape: torch.Size([2, 2, 4])

reshape changes the shape of a tensor without changing its data. The new shape must have the same number of elements as the original tensor. When using reshape, the memory is not guaranteed to be contiguous, so it may return a copy of the data.

As shown in above example, we can also use -1 to tell PyTorch to infer the missing dimension.

On the other hand, view always shares memory with the original tensor (a zero-copy operation). It requires the tensor to be contiguous in memory. If it’s not, it raises an error.

# view shares memory with the original tensor
x[0] = 100
viewed[1, 1, 3] = 1000
x, viewed

### Output
(tensor([ 100,    1,    2,    3,    4,    5,    6,    7,    8,    9,   10,   11,
          12,   13,   14, 1000]),
 tensor([[[ 100,    1,    2,    3],
          [   4,    5,    6,    7]],

         [[   8,    9,   10,   11],
          [  12,   13,   14, 1000]]]))

The important observation is that modifying viewed also modifies x, it only creates a view of the original tensor’s underlying storage rather than necessarily creating an independent copy. This could be very useful in memory managment, as long as you understand how it works of course.

torch.stackand torch.cat are two other important operations for combining tensors. stack creates a new dimension, while cat concatenates along an existing dimension.

a = torch.rand(2, 4)
b = torch.rand(2, 4)

stacked_0 = torch.stack([a, b], dim=0)
stacked_1 = torch.stack([a, b], dim=1)
cat_0 = torch.cat([a, b], dim=0)

print('shape of a:', a.shape)
print('shape of b:', b.shape)
print('shape of stacked_0:', stacked_0.shape)
print('shape of stacked_1:', stacked_1.shape)
print('shape of cat_0:', cat_0.shape)

### Output
shape of a: torch.Size([2, 4])
shape of b: torch.Size([2, 4])
shape of stacked_0: torch.Size([2, 2, 4])
shape of stacked_1: torch.Size([2, 2, 4])
shape of cat_0: torch.Size([4, 4])

torch.stack() creates takes two tensors a and b (each shape (2, 4)) and stacks them along a new first dimension. This mean dimension 0 now selects which original tensor you want,

stacked_0[0] is a
stacked_0[1] is b

torch.cat() joins tensors along an existing dimeension. That is it take the tensors a and b (each shape (2, 4)) and concat them along along dimension 0 gives a new tensor of shape (4, 4). The first two rows are from a and the last two rows are from b.

A useful mental model is:

stack creates a new dimension; cat joins along an existing dimension.

torch.squeeze and torch.unsqueeze manipulate dimensions of size 1.

x_reshaped = a.reshape(1, 2, 4, 1)
x_squeezed = x_reshaped.squeeze()

x_reshaped.shape, x_squeezed.shape, x_squeezed.unsqueeze(0).shape, x_squeezed.unsqueeze(1).shape

### Output
(torch.Size([1, 2, 4, 1]),
 torch.Size([2, 4]),
 torch.Size([1, 2, 4]),
 torch.Size([2, 1, 4]))

squeeze() removes dimensions whose size is 1 whereas unsqueeze() does the opposite: it inserts a new dimension of size 1 at a given index.

This is particularly useful when adding a batch dimension to a single input.

permute() changes the order of existing dimensions in the specified order.

x_reshaped.shape, x_reshaped.permute(0, 3, 1, 2).shape, x_reshaped.permute(2, 1, 3, 0).shape

### Output
(torch.Size([1, 2, 4, 1]), torch.Size([1, 1, 2, 4]), torch.Size([4, 2, 1, 1]))

This is different from reshape(), which changes how the elements are organised into dimensions. On the other handpermute() changes the order of the existing dimensions.

⋅ • ✦ • ⋅

Tensor Indexing

Indexing allows us to access particular elements or slices of a tensor.

my_tensor = torch.rand([2, 2, 3])

print("full tensor:\n", my_tensor)
print("my_tensor[1]:\n", my_tensor[1])
print("my_tensor[1,1]:\n", my_tensor[1, 1])
print("my_tensor[1,1,1]:", my_tensor[1, 1, 1])
print("my_tensor[:,1,1]:", my_tensor[:, 1, 1])

### Output
full tensor:
 tensor([[[0.6061, 0.0931, 0.1111],
         [0.6788, 0.6335, 0.0631]],

        [[0.0011, 0.4727, 0.3266],
         [0.5408, 0.6700, 0.9118]]])
my_tensor[1]:
 tensor([[0.0011, 0.4727, 0.3266],
        [0.5408, 0.6700, 0.9118]])
my_tensor[1,1]:
 tensor([0.5408, 0.6700, 0.9118])
my_tensor[1,1,1]: tensor(0.6700)
my_tensor[:,1,1]: tensor([0.6335, 0.6700])

my_tensor[1] Selects the second element of dimension 0.

my_tensor[1, 1] Selects the second element of dimensions 0 and 1.

my_tensor[1, 1, 1] Selects one exact element.

my_tensor[:, 1, 1] Selects all elements in dimension 0 at index 1 of dimension 1 and index 1 of dimension 2.

The colon means, take every element along the specified dimension.

⋅ • ✦ • ⋅

Moving Tensors Between Devices

When using MPS, you will often need to move tensors between the CPU and GPU.

PyTorch provides the .to() method for this.

device = "mps" if torch.backends.mps.is_available() else "cpu"
x = torch.ones([2, 2], device=device)
y = torch.ones([2, 2])

print("x device:", x.device)
print("y device before to():", y.device)
print("y device after to(device):", y.to(device).device)
print("x moved to cpu device:", x.to("cpu").device)

Output on an Apple Silicon Mac with MPS

### Output
x device: mps:0
y device before to(): cpu
y device after to(device): mps:0
x moved to cpu device: cpu

Notice that y initially lives on the CPU because no device was specified:

y = torch.ones([2, 2])

We can then create a copy on the selected device with:

y.to(device)

Likewise, an MPS tensor can be moved back to the CPU using:

x.to("cpu")

Why this is important is that PyTorch operations generally require all tensors to be on the same device. If you try to perform an operation between tensors on different devices, you will get an error.

Suppose you have:

model → MPS
input → CPU

You generally cannot simply perform the model operation using those tensors because they are located on different devices.

Instead, you typically move the input to the same device as the model:

CPU input
    ↓
.to("mps")
    ↓
MPS input
    ↓
MPS model

This is one of the most common patterns when training neural networks on Apple Silicon.

Another important consideration is that moving data between CPU and GPU takes time.

Therefore, you generally want to avoid repeatedly doing:

CPU → MPS → CPU → MPS → CPU

inside a performance-critical loop.

Instead, keep the model and the majority of the computation on the same device for as long as practical.

Wrap-up

This is by no means a complete list of the basic operations available in PyTorch. Having worked extensively with NumPy and Pandas, I find PyTorch surprisingly intuitive ,and honestly, quite fun to use.

I hope this tutorial has given you a practical understanding of some of the most important tensor operations you’ll encounter when building and training neural networks:

  • tensor creation
  • data types and devices
  • arithmetic and matrix algebra
  • shape transformations and indexing
  • CPU versus MPS execution
  • moving tensors between devices

I’ve spent quite a bit of time debugging PyTorch code only to discover that the problem was something incredibly simple. 😄 Over time, I’ve found it useful to go back to the basics and ask three simple questions:

Is the shape correct? — probably the most important one.

Is the dtype correct? — especially important when working with neural networks.

Is the device correct? — CPU, MPS, or something else?

Many PyTorch errors and plenty of unexpected results ultimately come down to one of these three properties. When things don’t behave as expected, checking shape, dtype, and device first can save you a surprising amount of debugging time.