Table of Contents
Moving Tensors and Models to the Right Device
In PyTorch, every tensor lives on a device, most commonly the CPU or a GPU. A model is just a collection of parameters, and each parameter is a tensor that also lives on a device. The first practical rule for training loops is that the model parameters and every tensor used in the forward pass must be on the same device, otherwise PyTorch will raise a runtime error.
A typical pattern is to pick a device once, move the model to it, then ensure each batch is moved to that same device right before the forward pass. In code, this usually means defining device = torch.device("cuda") when CUDA is available, otherwise falling back to torch.device("cpu"), then calling model.to(device). For each batch, you call x = x.to(device) and y = y.to(device).
All tensors that interact in an operation must be on the same device. If your model is on CUDA, your inputs and targets must be moved to CUDA before the forward pass.
When moving batches, you will often see x = x.to(device, non_blocking=True) when using a DataLoader with pinned memory. That combination can reduce input transfer overhead when training on GPU, but it is an optimization, not a requirement.
Checking and Debugging Device Issues
Device mistakes are among the most common beginner issues because they can hide until the first real batch hits a layer. It helps to verify early. You can check a tensor’s device via x.device and a model parameter device via next(model.parameters()).device. If those do not match, fix that first.
Another frequent source of confusion is creating new tensors inside the training step. For example, calling torch.zeros(...) without specifying a device creates a CPU tensor by default. If you need a new tensor during the forward pass or loss computation, create it on the same device as existing tensors, often by using torch.zeros_like(existing_tensor) or torch.zeros(..., device=x.device).
New tensors you create inside the training loop default to CPU unless you explicitly place them on the correct device. Prefer *_like helpers to automatically match device and dtype.
What Mixed Precision Is and Why It Helps
Mixed precision training uses a combination of floating point formats, typically FP16 or BF16 for many computations, while keeping some values in FP32 for numerical stability. The main benefits on modern GPUs are faster throughput and lower memory usage, which can allow larger batch sizes or bigger models.
Two concepts matter in practice. The first is autocasting, which tells PyTorch to run certain operations in lower precision when it is safe. The second is gradient scaling, which helps avoid underflow in FP16 gradients by multiplying the loss by a scale factor during backpropagation, then unscaling before the optimizer step.
BF16 is generally more numerically stable than FP16 because it has a larger exponent range. On supported hardware, BF16 often needs no gradient scaling. FP16 usually benefits from gradient scaling.
Using `torch.cuda.amp` in a Training Loop
The standard PyTorch approach for mixed precision on CUDA uses torch.autocast and torch.cuda.amp.GradScaler. You keep model parameters in FP32, and you wrap the forward and loss computation in an autocast context. Backward uses scaled loss, then you step the optimizer through the scaler.
A minimal training step looks like this:
scaler = torch.cuda.amp.GradScaler(enabled=(device.type == "cuda"))
for x, y in train_loader:
x = x.to(device, non_blocking=True)
y = y.to(device, non_blocking=True)
optimizer.zero_grad(set_to_none=True)
with torch.autocast(device_type="cuda", dtype=torch.float16, enabled=(device.type == "cuda")):
preds = model(x)
loss = loss_fn(preds, y)
scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()This pattern works because autocast affects forward computations, while GradScaler manages the backward and step safely for FP16. If you are not on CUDA, you typically disable autocast and scaling so the same code path still runs.
When using FP16 mixed precision, do not call optimizer.step() directly. Use scaler.step(optimizer) and scaler.update() so that gradients are properly unscaled and invalid steps can be skipped.
Mixed Precision on CPU and Apple Silicon
Mixed precision is primarily beneficial on GPUs. Autocast also exists for CPU, but the speedups are usually smaller and depend heavily on hardware and workload. On Apple Silicon, PyTorch uses the MPS backend for GPU acceleration, and mixed precision support and behavior can differ from CUDA. For absolute beginners, the simplest, most reliable approach is to use mixed precision only when training on CUDA, and to keep CPU and MPS training in full precision unless you have a reason and have verified stability.
Common Stability Problems and How to React
If your loss becomes NaN or training diverges after enabling mixed precision, first confirm that your non mixed precision run is stable. Then try switching FP16 to BF16 if your GPU supports it, because BF16 is often more forgiving. Another safe fallback is to disable mixed precision entirely for that experiment. Mixed precision is an optimization, not a requirement.
Also be careful with operations known to be numerically sensitive, such as certain reductions, exponentials, or softmax in custom code. Autocast handles many cases well, but if you write custom math, you may need to compute sensitive parts in FP32.
Mixed precision should not change your model’s correctness goals. If enabling it makes training unstable, switch to BF16 if available or disable mixed precision and continue with FP32.
Practical Rules of Thumb for Beginners
Pick one device variable and use it everywhere, move the model once, and move each batch right before the forward pass. Avoid creating stray CPU tensors inside the training step. Enable CUDA mixed precision only after your baseline training loop works, then adopt the standard autocast plus GradScaler pattern. Keep the code structured so you can flip mixed precision on or off with a single flag, which makes debugging much easier.