PyTorch GPU/CUDA

Practical GPU usage in PyTorch — moving models and data to the GPU correctly, and the common device-related errors that come up in real training code.

The Standard Device Pattern

import torch

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(device)   # cuda:0 or cpu, depending on what's available

model = MyModel().to(device)   # moves ALL of the model's parameters to the device

for x_batch, y_batch in train_loader:
    x_batch = x_batch.to(device)
    y_batch = y_batch.to(device)
    predictions = model(x_batch)   # now everything is on the same device

Why the "Same Device" Rule Exists

Every tensor operation (matrix multiplication, addition) requires its operands to physically reside in the same memory space — a CPU tensor and a GPU tensor can't be combined directly, since they exist in entirely separate physical memory. This is precisely why the classic "Expected all tensors to be on the same device" error occurs whenever a model and its input data end up on different devices.

Moving Data Back to CPU (for NumPy, Plotting, etc.)

gpu_tensor = torch.randn(3, 3).to(device)

cpu_tensor = gpu_tensor.cpu()          # move back to CPU
numpy_array = gpu_tensor.cpu().numpy()   # .numpy() ONLY works on CPU tensors -- must move first
value = gpu_tensor.item()               # extracting a single scalar also requires it to be accessible

Mixed Precision Training — A Brief Practical Note

from torch.cuda.amp import autocast, GradScaler

scaler = GradScaler()

for x_batch, y_batch in train_loader:
    optimizer.zero_grad()
    with autocast():   # automatically uses lower precision (FP16) where numerically safe
        predictions = model(x_batch)
        loss = loss_fn(predictions, y_batch)
    scaler.scale(loss).backward()   # scales the loss to prevent FP16 gradient underflow
    scaler.step(optimizer)
    scaler.update()

Mixed precision training uses lower-precision (FP16) computation where it's numerically safe, speeding up training and reducing memory usage on modern GPUs, while GradScaler prevents small gradient values from underflowing to zero in the lower-precision format.

Common Mistakes

• Calling .numpy() directly on a GPU tensor — this raises an error; the tensor must first be moved to CPU via .cpu().
• Moving only the model or only the data to the GPU, forgetting the other — both must be explicitly moved for the device-mismatch error to be avoided.
• Checking torch.cuda.is_available() once at import time and hardcoding the result, rather than using it to set a device variable consistently referenced throughout the code — makes code far less portable between GPU and CPU-only environments.

Interview Relevance

Q: "Why do you need to call .to(device) on both the model and every batch of input data, rather than just one or the other?" Every tensor operation requires its operands to physically reside in the same memory space — a model's weights and the input data it processes must both be on the same device (CPU or a specific GPU) for the underlying computation (matrix multiplication, addition) to execute at all; mismatched devices raise a runtime error rather than silently working across memory spaces.

Practice Question

Why does GradScaler matter specifically for mixed precision (FP16) training, but not for standard FP32 training?

Related DL Notes

• PyTorch Training Loop
• GPU Utilization
• Memory Optimization