TopolearnGPUs make deep learning 10-100x faster, but only if data and model actually live there.
To train on a GPU, put both the model and the data on it:
model.to(device)moves all parameters and buffers in place. Do this before creating the optimizer.- Each batch:
xb = xb.to(device, non_blocking=True). Withpin_memory=Truein the DataLoader, the copy can overlap with computation.
Write the device choice once and keep the rest device-agnostic. CUDA covers NVIDIA (and AMD via ROCm builds, which also use the "cuda" name); "mps" covers Apple silicon. Recent releases also offer torch.accelerator.current_accelerator(), which picks whichever accelerator is present.
GPU work is asynchronous: Python queues kernels and continues. The CPU only waits when it needs a result, for example .item(), .cpu(), print(tensor) or an if on a tensor's value. Each of these is a synchronisation; too many in the inner loop starve the GPU. When timing GPU code, call torch.cuda.synchronize() before reading the clock.
Check memory with torch.cuda.memory_allocated() and nvidia-smi. If you run out of memory, reduce the batch size first, then look at mixed precision and gradient accumulation.
Common hiccups: tensors created inside forward with torch.zeros(...) land on the CPU; create them with device=x.device or torch.zeros_like.
Example
import torch
from torch import nn
device = torch.device(
"cuda" if torch.cuda.is_available()
else "mps" if torch.backends.mps.is_available()
else "cpu")
model = nn.Linear(10, 2).to(device) # move parameters first
opt = torch.optim.SGD(model.parameters(), lr=0.1)
x = torch.randn(64, 10).to(device, non_blocking=True)
mask = torch.ones(64, 1, device=x.device) # create new tensors on the same device
loss = (model(x) * mask).pow(2).mean()
opt.zero_grad()
loss.backward()
opt.step()
print(loss.device.type, round(loss.item(), 3)) # .item() synchronises and copies to CPU
if device.type == "cuda":
torch.cuda.synchronize() # before timing anything
print(torch.cuda.memory_allocated() // 2**20, "MiB allocated")Remember
- Model and data must both be on the device
- GPU calls are async;
.item(),.cpu()and printing synchronise - Create tensors inside
forwardwithdevice=x.device
Watch out for
- Moving the inputs to the GPU but forgetting the targets, which fails inside the loss.