Devices and Tensor Placement

The CPU/GPU border is a real thing

PyTorch tensors live on a specific device. x.device tells you which one. Operations between tensors require both to be on the same device, and you have to move them explicitly with .to(device). Forgetting to do this is the most common error message new PyTorch users see: Expected all tensors to be on the same device.

The model and the data both have to be on the GPU for training to be fast. Move the model once at startup; move each batch in the training loop. Output tensors stay on the GPU until you explicitly bring them back (e.g. for printing or for non-PyTorch downstream code).

Tip: Three places to put .to(device): once on the model at startup, once on each input batch in the loop, and once on the loss tensor when you want to print it (use .item() to bring a scalar back to Python). Almost everything else stays on the GPU.

Picking the right device

One liner that handles the three common cases (NVIDIA, Apple Silicon, CPU fallback) — see code below. Apple's MPS backend is fast for many workloads but still has occasional missing ops (look for NotImplementedError and report them upstream).

Multi-GPU and distributed

For one machine, multiple GPUs: use torch.nn.parallel.DistributedDataParallel (DDP), launched with torchrun. For multiple machines: same DDP API, different launch script. Higher-level: accelerate launch from Hugging Face hides the boilerplate.

Principle: In 2026, default to single-GPU training until your model genuinely does not fit. The complexity tax of multi-GPU is real, and modern accelerators (H100, MI300X, M3 Ultra) hold a surprising amount of model.

Code

Pick the device once, use it everywhere·python

import torch

device = (
    torch.device("cuda") if torch.cuda.is_available()
    else torch.device("mps") if torch.backends.mps.is_available()
    else torch.device("cpu")
)
print("device:", device)

model = MyModel().to(device)

for xb, yb in train_loader:
    xb, yb = xb.to(device, non_blocking=True), yb.to(device, non_blocking=True)
    logits = model(xb)
    loss = loss_fn(logits, yb)
    loss.backward()
    optimizer.step()
    print(loss.item())   # .item() brings the scalar back to CPU/Python

The CPU/GPU border is a real thing

Picking the right device

Multi-GPU and distributed

Code

External links

Exercise

Progress

Comments 0