04 · Device Placement 🖥️⚡

4 min read

This is the sub-module that makes PyTorch fast — and the source of the single most common error beginners hit.

1. The 'Why'

A CPU is a brilliant generalist: it handles a handful of complex tasks quickly. A GPU is a specialist: it does thousands of simple arithmetic operations simultaneously. Neural networks are basically enormous piles of simple arithmetic (all those matrix multiplies from the previous sub-module), which is exactly the workload GPUs were born for. The practical consequence is enormous — a model that trains in an hour on a GPU might take a week on a CPU. Knowing how to place your tensors on the right hardware is therefore not an optional nicety; it's the difference between iterating on ideas and staring at a progress bar.

The catch is that PyTorch requires the tensors involved in an operation to live on the same device. Mix a CPU tensor with a GPU tensor and you get an immediate RuntimeError. This trips up nearly every beginner the first time they move a model to a GPU but forget to move the data too. This sub-module gives you a clean, portable pattern for handling devices — including Apple Silicon's MPS backend and NVIDIA's CUDA — so you never have to think about it twice.

2. Core Concepts

The device attribute. Every tensor knows where it lives: cpu, cuda (NVIDIA GPU), or mps (Apple Silicon GPU). Check it with t.device.

Moving tensors. t.to(device) returns a new tensor on the target device (the original is unchanged). You can also create tensors directly on a device with device= at construction time, which avoids an extra copy.

The same-device rule. Both operands of an operation must be on the same device. The fix is a discipline, not a trick: define one device variable up front and move both your model and every batch of data onto it.

The NumPy bridge. Tensors convert to/from NumPy with .numpy() and torch.from_numpy(...), but only on the CPU — you must .cpu() a GPU tensor first. Note that from_numpy shares memory with the NumPy array.

3. Code in Action

Portable device selection

import torch

# Pick the best available device without hard-coding "cuda"
if torch.cuda.is_available():
    device = torch.device("cuda")       # NVIDIA GPU
elif torch.backends.mps.is_available():
    device = torch.device("mps")        # Apple Silicon GPU
else:
    device = torch.device("cpu")        # fallback: always works
print("Using:", device)

Moving and creating tensors on a device

import torch
device = "cuda" if torch.cuda.is_available() else "cpu"

x = torch.rand(3, 3)         # created on the CPU by default
x = x.to(device)             # move to the chosen device (returns a NEW tensor)
print(x.device)              # reflects the destination device

y = torch.ones(3, 3, device=device)  # create DIRECTLY on the device (no copy)
z = x + y                             # works: both operands share a device

The classic error — and the NumPy bridge

import torch
import numpy as np

a = torch.ones(3)                 # on CPU
b = torch.ones(3, device="cuda") if torch.cuda.is_available() else torch.ones(3)
# a + b would raise: "Expected all tensors to be on the same device"
# Fix by aligning devices first:
result = a.to(b.device) + b       # move `a` to match `b`, then add

# NumPy conversion must happen on the CPU
np_array = result.cpu().numpy()   # .cpu() first, then .numpy()
back = torch.from_numpy(np_array) # NumPy -> tensor (shares memory!)

4. Common Pitfalls

Pitfall 1 — "Expected all tensors to be on the same device." The most common beginner error: the model is on the GPU but a batch of data is still on the CPU. Fix: define one device variable and call .to(device) on both the model and every input batch. Never mix.

Pitfall 2 — Calling .numpy() on a GPU tensor. This throws an error because NumPy is CPU-only. Always .cpu() first: tensor.cpu().numpy(). Also remember from_numpy shares memory, so mutating the array mutates the tensor — .clone() if you need independence.

Pitfall 3 — Assuming .to() mutates in place. For tensors, x.to(device) returns a new tensor and leaves the original where it was, so x.to(device) alone does nothing useful — you must reassign: x = x.to(device). (Models are the exception: model.to(device) moves parameters in place.)

5. Further Reading & Watch List

⬅️ Prev: 03 · Tensor Math · 🏁 Module complete! Back to the module index or on to Module 02 — Autograd.