A model sitting on the GPU and a batch still on the CPU: that one mismatch kills more first runs than any piece of math. PyTorch GPU training is mostly two habits, keeping everything on the same device and reading error messages instead of panicking at them. This post covers device discipline, CUDA out of memory, shape bugs, NaN loss, and mixed precision speedups.
“Neural nets fail silently. The code runs, the loss looks plausible, and the bug just quietly costs you accuracy. Learn to be suspicious.”
Andrej Karpathy, “A Recipe for Training Neural Networks”
Last Updated: July 2026 | Tested on: Python 3.14.6, PyTorch 2.12.1 | Difficulty: Advanced | Reading Time: 23 minutes
A GPU (Graphics Processing Unit) is not magic. It is a second, separate pool of memory (called VRAM) sitting next to your normal system RAM, packed with thousands of small cores that do the same math on many numbers at once. That parallelism is why a matrix multiply that crawls on a Central Processing Unit (CPU) flies on a GPU. The catch is the “separate memory” part: your data starts life in system RAM, your model is born in system RAM, and nothing runs on the GPU until you deliberately copy it across.
Forget to copy one of them and PyTorch stops you with a device-mismatch error. Most of this post is teaching you the moves that keep both sides in sync, then the debugging reflexes for when they are not.
Here is what we cover:
- The .to(device) discipline and the device-mismatch RuntimeError
- Reading and fixing a CUDA out of memory error
- Debugging shape mismatches and NaN loss
- Printing the gradient norm to catch exploding or vanishing gradients
- Mixed precision training for a real speedup
- Distributed training literacy: DDP, FSDP, and friends
- Reproducibility with seeds, and where seeds stop helping
Table of Contents
Prerequisites
Picture two kitchens in one house. The pantry (system RAM) is where all your groceries live, and the busy prep counter (VRAM) is where the actual cooking happens fast. Food does not teleport from pantry to counter; someone has to carry it. In PyTorch you are that someone, and .to(device) is the act of carrying. The diagram below shows what travels where during training, and how often.
Tip: click the diagram to open it full screen, then use the zoom and pan controls for a closer look.
Notice the frequencies. The model crosses to VRAM once at the start and then stays put; the optimizer updates it in place, so you never copy weights back and forth mid-training. Each batch of data, on the other hand, gets copied every single iteration, which is why an efficient input pipeline matters so much. And when you log the loss, only one tiny number travels back to the CPU, which is exactly why you write loss.item() and not loss.
- First neural network with PyTorch (tensors, autograd, the training loop)
- A machine with an NVIDIA GPU, or a free Google Colab GPU runtime, to see the speedups
- pip install torch torchvision
One honest note. The debugging code here (shape errors, NaN loss, gradient norms, seeds) is tested on a plain CPU, so you can run every line on any laptop and get the exact numbers shown. The two blocks that only make sense on real hardware, the CUDA out of memory message and the timed speedup, are clearly labeled as example output from a GPU runtime.
The .to(device) Discipline
The whole PyTorch GPU game is one rule: the model and the data it touches must be on the same device. Pick the device once at the top of your script, then move everything to it. Never hardcode "cuda", because the same script should still run on a CPU-only laptop or an Apple Silicon Mac. The pattern below checks what is available and picks the best option.
📄 device_discipline.py: pick one device, move both model and data
import torch
# Pick the best device once, at the top of your script
if torch.cuda.is_available():
device = torch.device("cuda")
elif torch.backends.mps.is_available(): # Apple Silicon
device = torch.device("mps")
else:
device = torch.device("cpu")
print(f"Training on: {device}")
# A tiny model and a batch of fake data
model = torch.nn.Linear(4, 3)
batch = torch.randn(8, 4)
# The discipline: move BOTH model and data to the same device
model = model.to(device)
batch = batch.to(device)
out = model(batch)
print(f"Model device: {next(model.parameters()).device}")
print(f"Batch device: {batch.device}")
print(f"Output device: {out.device}, shape: {tuple(out.shape)}")
▶ Output (on a CPU-only machine)
Training on: cpu Model device: cpu Batch device: cpu Output device: cpu, shape: (8, 3)
What happened here: on our test machine there is no GPU, so the code falls through to cpu and everything lines up. On a Colab GPU runtime the first line would read Training on: cuda and every device below it would say cuda:0. The important part is that model and batch always match, because we sent both to the same device variable. One subtlety: model.to(device) moves the model in place, but for tensors .to(device) returns a new tensor, so you must reassign with batch = batch.to(device). Forgetting that reassignment is the number one reason people think they moved their data when they did not.
So what does it look like when you break the rule? Say a user named Aditi moves her model to the GPU but forgets to move the batch. PyTorch refuses to multiply a CPU tensor by a GPU weight and throws a very specific error.
▶ Example output (running on a GPU with a misplaced tensor)
RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu! (when checking argument for argument mat1 in method wrapper_CUDA_addmm)
What happened here: the message names both devices it found, cuda:0 and cpu, and even tells you which argument was wrong (mat1, the input to a linear layer). The fix is never clever: find the tensor that is still on the CPU and add .to(device) to it. This usually bites on things you create mid-loop, like a fresh tensor of class weights or a positional index, that you forgot to move. Read the two device names, move the stray one, done.
Reading a CUDA Out of Memory Error
VRAM is small and precious. A mid-range GPU might have 8 to 16 GB, and your model, its gradients, the optimizer state, and every activation from the forward pass all have to fit. Ask for too much and you get the error every PyTorch GPU practitioner meets in their first week.
▶ Example output (from a GPU runtime under memory pressure)
torch.OutOfMemoryError: CUDA out of memory. Tried to allocate 2.00 GiB. GPU 0 has a total capacity of 15.77 GiB of which 1.43 GiB is free. Process 0 has 14.34 GiB memory in use. Of the allocated memory 13.10 GiB is allocated by PyTorch, and 210.00 MiB is reserved but unallocated. If reserved but unallocated memory is large try setting PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
What happened here: read it top to bottom. It wanted 2 GiB, only 1.43 GiB was free, so it stopped. The single most effective fix is to lower the batch size, because activation memory scales almost linearly with it. Halving the batch roughly halves activation memory. But smaller batches can make training noisier, so the trick is to keep the same effective batch size using gradient accumulation: run several small batches, let their gradients pile up, and step the optimizer only once. The next block does exactly that, and it runs fine on a CPU so you can see the mechanics.
📄 grad_accumulation.py: an effective batch of 64 while only holding 16 at a time
import torch
import torch.nn as nn
torch.manual_seed(3)
X = torch.randn(64, 10)
y = torch.randint(0, 3, (64,))
crit = nn.CrossEntropyLoss()
# Goal: an effective batch of 64, but pretend VRAM only fits 16 at a time.
micro_batch = 16
accum_steps = 4 # 16 * 4 = 64 effective
model = nn.Sequential(nn.Linear(10, 32), nn.ReLU(), nn.Linear(32, 3))
opt = torch.optim.SGD(model.parameters(), lr=0.1)
opt.zero_grad()
for i in range(accum_steps):
xb = X[i * micro_batch:(i + 1) * micro_batch]
yb = y[i * micro_batch:(i + 1) * micro_batch]
loss = crit(model(xb), yb) / accum_steps # scale so the sum matches a full batch
loss.backward() # gradients ACCUMULATE, no step yet
print(f" micro-step {i}: partial loss = {loss.item():.4f}")
opt.step() # one real update using all 64 samples worth of gradient
opt.zero_grad()
print("Stepped once on an effective batch of 64 while only holding 16 in memory.")
▶ Output
micro-step 0: partial loss = 0.2880 micro-step 1: partial loss = 0.2820 micro-step 2: partial loss = 0.2859 micro-step 3: partial loss = 0.2715 Stepped once on an effective batch of 64 while only holding 16 in memory.
What happened here: we never called zero_grad() inside the loop, so each backward() added its gradients on top of the previous one. After four micro-batches the .grad buffers hold the summed gradient of all 64 samples, and a single step() applies it. The one detail people miss is dividing the loss by accum_steps: without it you would be summing four full-size gradients and effectively quadrupling your learning rate. Beyond accumulation, other memory savers are mixed precision (covered below), gradient checkpointing (recompute activations instead of storing them), and simply choosing a smaller model.
Debugging Clinic: Shape Mismatches
Shape errors are the most common bug in all of deep learning, and they are also the friendliest, because PyTorch tells you the exact dimensions that did not line up. A linear layer built for 4 input features cannot multiply a row of 5 numbers. Here is that mistake and its fix.
📄 shape_bug.py: trigger and read a shape mismatch
import torch
import torch.nn as nn
torch.manual_seed(0)
# A model that expects 4 input features
model = nn.Sequential(
nn.Linear(4, 16),
nn.ReLU(),
nn.Linear(16, 3),
)
# But Aditi accidentally feeds it 5 features per row
bad_batch = torch.randn(8, 5)
try:
out = model(bad_batch)
except RuntimeError as e:
print("RuntimeError:")
print(e)
# The fix: match the input width the first layer was built for
good_batch = torch.randn(8, 4)
out = model(good_batch)
print(f"\nFixed. Output shape: {tuple(out.shape)}")
▶ Output
RuntimeError: mat1 and mat2 shapes cannot be multiplied (8x5 and 4x16) Fixed. Output shape: (8, 3)
What happened here: the message reads (8x5 and 4x16). That is your batch (8 rows, 5 features) trying to multiply the first layer’s weight matrix (4 by 16). The inner numbers, 5 and 4, must match for a matrix multiply, and they do not. Once you learn to read that pair, shape bugs become a ten-second fix. When the mismatch is deeper in the network, drop a print(x.shape) between layers, or use a single known batch and step through with a debugger. Convolutions and attention layers have their own shape rules, but the debugging move is identical: print shapes, compare the numbers, find where they stop matching.
Debugging Clinic: NaN Loss
A NaN (Not a Number) loss is training’s version of a dropped call: everything looks connected, but nothing is getting through anymore. Once the loss becomes NaN, every gradient becomes NaN, every weight becomes NaN, and the model is dead. The most common cause is a learning rate that is too high, which makes each update overshoot, so the loss grows, which makes the next update bigger, until the numbers overflow. Let us watch it happen.
📄 nan_loss.py: a too-high learning rate blows the loss up to NaN
import torch
import torch.nn as nn
import math
torch.manual_seed(1)
X = torch.randn(200, 10)
y = torch.randint(0, 3, (200,))
def train(lr, epochs=8):
torch.manual_seed(1)
model = nn.Sequential(nn.Linear(10, 32), nn.ReLU(), nn.Linear(32, 3))
opt = torch.optim.SGD(model.parameters(), lr=lr)
crit = nn.CrossEntropyLoss()
losses = []
for _ in range(epochs):
opt.zero_grad()
loss = crit(model(X), y)
loss.backward()
opt.step()
losses.append(loss.item())
return losses
print("Learning rate 100.0 (way too high):")
for i, l in enumerate(train(100.0, epochs=20)):
if math.isnan(l):
print(f" epoch {i}: loss = nan <-- training is dead, gradients gone")
break
if math.isinf(l):
print(f" epoch {i}: loss = inf <-- overflowed float32, nan is next")
else:
print(f" epoch {i}: loss = {l:.4g}")
print("\nLearning rate 0.1 (sane):")
for i, l in enumerate(train(0.1, epochs=6)):
print(f" epoch {i}: loss = {l:.4f}")
▶ Output
Learning rate 100.0 (way too high): epoch 0: loss = 1.118 epoch 1: loss = 6.851 epoch 2: loss = 4589 epoch 3: loss = 9.248e+05 epoch 4: loss = 2.235e+08 epoch 5: loss = 6.15e+10 epoch 6: loss = 1.875e+13 epoch 7: loss = 7.435e+15 epoch 8: loss = 1.752e+18 epoch 9: loss = 6.712e+20 epoch 10: loss = 1.643e+23 epoch 11: loss = 6.177e+25 epoch 12: loss = 1.67e+28 epoch 13: loss = 6.026e+30 epoch 14: loss = 1.429e+33 epoch 15: loss = 2.823e+35 epoch 16: loss = nan <-- training is dead, gradients gone Learning rate 0.1 (sane): epoch 0: loss = 1.1185 epoch 1: loss = 1.1166 epoch 2: loss = 1.1147 epoch 3: loss = 1.1130 epoch 4: loss = 1.1113 epoch 5: loss = 1.1096
What happened here: with the learning rate at 100 the loss roughly multiplies each epoch, 4589, then 900 thousand, then hundreds of millions, until at epoch 15 it passes the largest number a 32-bit float can hold (about 3.4e38), overflows to infinity, and the very next step turns infinity into NaN. Drop the rate to 0.1 and the same model trains calmly.
When you hit a NaN, walk this checklist in order: lower the learning rate first, then check your inputs for NaN or infinity (a stray divide-by-zero in preprocessing does it), then look for a log(0) or sqrt of a negative in a custom loss, and add gradient clipping, which we cover next. Nine times out of ten it is the learning rate.
Watching Gradients: Exploding and Vanishing
You do not have to guess whether gradients are misbehaving. You can measure them. The gradient norm is a single number that says how big the whole update is: add up the square of every gradient, take the square root. A norm that climbs toward huge values means exploding gradients; a norm that shrinks toward zero means vanishing gradients, where deep layers stop learning. Print it every few steps and you have a dashboard for training health. The classic fix for the exploding case is clip_grad_norm_, which rescales the gradient so its norm never exceeds a ceiling you set.
📄 grad_norm.py: measure the gradient norm and clip it
import torch
import torch.nn as nn
torch.manual_seed(2)
X = torch.randn(64, 10)
y = torch.randint(0, 3, (64,))
model = nn.Sequential(nn.Linear(10, 64), nn.ReLU(), nn.Linear(64, 3))
opt = torch.optim.SGD(model.parameters(), lr=0.5)
crit = nn.CrossEntropyLoss()
def grad_norm(m):
total = 0.0
for p in m.parameters():
if p.grad is not None:
total += p.grad.data.norm(2).item() ** 2
return total ** 0.5
print("Without clipping:")
for step in range(4):
opt.zero_grad()
loss = crit(model(X), y)
loss.backward()
print(f" step {step}: grad_norm = {grad_norm(model):.3f}")
opt.step()
# Reset and clip
torch.manual_seed(2)
model = nn.Sequential(nn.Linear(10, 64), nn.ReLU(), nn.Linear(64, 3))
opt = torch.optim.SGD(model.parameters(), lr=0.5)
print("\nWith clip_grad_norm_(max_norm=1.0):")
for step in range(4):
opt.zero_grad()
loss = crit(model(X), y)
loss.backward()
before = grad_norm(model)
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
after = grad_norm(model)
print(f" step {step}: {before:.3f} -> clipped to {after:.3f}")
opt.step()
▶ Output
Without clipping: step 0: grad_norm = 0.499 step 1: grad_norm = 0.268 step 2: grad_norm = 0.236 step 3: grad_norm = 0.219 With clip_grad_norm_(max_norm=1.0): step 0: 0.387 -> clipped to 0.387 step 1: 0.241 -> clipped to 0.241 step 2: 0.215 -> clipped to 0.215 step 3: 0.202 -> clipped to 0.202
What happened here: this small, well-behaved model has gradient norms well under 1.0, so clipping at max_norm=1.0 leaves them untouched, which is exactly what you want. Clipping is a safety net, not a straitjacket: it only kicks in when a norm exceeds the ceiling, and then it scales the whole gradient down proportionally. In a network prone to spikes, like an Recurrent Neural Network (RNN) or a Transformer, you would see norms occasionally jump to 50 or 100 and get pulled back to 1.0. If instead your norms drift toward 0.001 and stall, that is vanishing gradients, and the answer is architectural: residual connections, normalization layers, and activations like ReLU or GELU rather than sigmoid in deep stacks.
Mixed Precision Training for Real Speedups
By default PyTorch stores numbers as 32-bit floats. Modern GPUs, though, have special cores that crunch 16-bit floats much faster and use half the memory, which is where PyTorch GPU training picks up its biggest easy win. Mixed precision means running the heavy matrix math in 16-bit while keeping a few sensitive parts in 32-bit for stability. You wrap the forward pass in torch.autocast and PyTorch decides, per operation, which precision is safe. Here is the forward half, which runs on any device so you can see the dtypes change.
📄 autocast_demo.py: run the forward pass in lower precision
import torch
import torch.nn as nn
torch.manual_seed(4)
model = nn.Sequential(nn.Linear(64, 128), nn.ReLU(), nn.Linear(128, 10))
x = torch.randn(32, 64)
y = torch.randint(0, 10, (32,))
crit = nn.CrossEntropyLoss()
# Under autocast, PyTorch runs eligible ops in a lower-precision float
with torch.autocast(device_type="cpu", dtype=torch.bfloat16):
out = model(x)
loss = crit(out, y)
print(f"Output dtype inside autocast: {out.dtype}")
print(f"Loss dtype: {loss.dtype}")
print(f"Loss value: {loss.item():.4f}")
# Same forward pass in full precision for comparison
out_fp32 = model(x)
print(f"Output dtype outside autocast: {out_fp32.dtype}")
▶ Output
Output dtype inside autocast: torch.bfloat16 Loss dtype: torch.float32 Loss value: 2.2765 Output dtype outside autocast: torch.float32
What happened here: inside the autocast block the linear layers produced bfloat16 output, but notice the loss came back as float32. That is autocast being smart: matrix multiplies run in 16-bit for speed, while the loss, which needs precision, stays in 32-bit. On a real GPU you also wrap the backward pass with a GradScaler, which multiplies the loss by a big factor before backward() so tiny 16-bit gradients do not round down to zero, then unscales before the optimizer step. The payoff is measured on a GPU, not a CPU, so here is a representative timing from a Colab run of a small convolutional network.
▶ Example output (Colab T4 GPU, small CNN, one epoch)
Full precision (fp32): 18.4 s/epoch, peak VRAM 3.9 GB Mixed precision (amp): 9.7 s/epoch, peak VRAM 2.3 GB Speedup: 1.9x, memory saved: 41%
What happened here: roughly a 1.9x speedup and 40% less memory, for two extra lines of code. The exact number depends on your GPU and model; older cards without dedicated 16-bit cores see less benefit, and very small models see almost none because they are not compute-bound in the first place. As of PyTorch 2.12.1 the recommended entry points are torch.autocast and torch.amp.GradScaler. The older torch.cuda.amp spelling still works but is being phased out, so prefer the device-agnostic torch.amp names in new code.
Distributed Training, Explained
At some point one card is not enough, either because the model does not fit or because you want PyTorch GPU training to run faster across several of them. You do not need a cluster to understand the ideas, and interviewers love to check that you know the vocabulary, so here is the literacy version. Think of it like cooking a banquet: you can either give every cook the full recipe and split the guests between them, or, if the recipe itself is too big to hold, split the recipe across cooks who pass dishes down the line.
- DDP (DistributedDataParallel): the workhorse. Every GPU holds a full copy of the model and trains on a different slice of each batch. After each backward pass, the GPUs average their gradients (an "all-reduce") so all copies stay identical. This is data parallelism, and it is what you reach for first.
- FSDP (Fully Sharded Data Parallel): for models too big to fit on one GPU. It shards the parameters, gradients, and optimizer state across GPUs, gathering each layer's weights only for the moment they are needed. This is how you train models with tens of billions of parameters.
- DeepSpeed: a popular library from Microsoft whose ZeRO optimizer pioneered the sharding ideas FSDP later adopted. Named a lot in interviews for large-model training.
- NCCL (NVIDIA Collective Communications Library): the low-level plumbing that actually moves gradients between GPUs quickly. You rarely call it directly, but it is the backend DDP uses on NVIDIA hardware.
The mental model that ties them together: DDP copies the whole model to every GPU and splits the data; FSDP splits the model itself when it is too large to copy. You write the single-GPU code first, get it correct, and only then wrap it in DDP, because the training loop logic barely changes. That is the real reason we teach single-GPU discipline so thoroughly: distributed training is that same loop with a coordination layer on top.
Reproducibility: Seeds and Their Limits
Random weight initialization, shuffled batches, and dropout all depend on a random number generator. Seeding that generator makes a run repeatable, which matters when you are hunting a bug or comparing two ideas fairly. But there are three separate generators in play (Python's random, NumPy, and PyTorch), so you seed all of them. Here is the helper you should paste into every PyTorch GPU project.
📄 reproducibility.py: seed everything, then ask for determinism
import torch
import random
import numpy as np
def set_seed(seed=42):
random.seed(seed)
np.random.seed(seed)
torch.manual_seed(seed)
torch.cuda.manual_seed_all(seed) # no-op on CPU, needed on GPU
# Run 1
set_seed(42)
a = torch.randn(3)
# Run 2, same seed
set_seed(42)
b = torch.randn(3)
# Run 3, different seed
set_seed(99)
c = torch.randn(3)
print(f"seed 42, run 1: {a.tolist()}")
print(f"seed 42, run 2: {b.tolist()}")
print(f"seed 99, run 3: {c.tolist()}")
print(f"run 1 == run 2 ? {torch.equal(a, b)}")
# Ask for bit-exact determinism (may slow things down or raise if a kernel has no
# deterministic version)
torch.use_deterministic_algorithms(True, warn_only=True)
print("Deterministic algorithms requested (warn_only=True).")
▶ Output
seed 42, run 1: [0.33669036626815796, 0.12880940735340118, 0.23446236550807953] seed 42, run 2: [0.33669036626815796, 0.12880940735340118, 0.23446236550807953] seed 99, run 3: [0.6126858592033386, -1.1753536462783813, -0.7646492719650269] run 1 == run 2 ? True Deterministic algorithms requested (warn_only=True).
What happened here: the same seed gives bit-for-bit identical tensors, a different seed gives different ones. Now the limits, because this is where people get burned. Seeds make one machine repeat itself, but they do not guarantee the same numbers across different GPUs, different PyTorch or CUDA versions, or CPU versus GPU, because the underlying math kernels differ, and some GPU operations are nondeterministic by design for speed. torch.use_deterministic_algorithms(True) forces the deterministic versions, but it can be slower and will raise if an operation has no deterministic version, which is why warn_only=True is a gentler start. With a multi-worker DataLoader you also seed each worker for the shuffling to repeat. Reproducibility is a spectrum, not a switch.
Common Mistakes
- Moving the model but not the data (or vice versa): the classic device-mismatch RuntimeError. Send both to the same
device, and remember tensors need reassignment:x = x.to(device). - Accumulating loss with
total_loss += lossinstead ofloss.item(): that keeps every batch's computation graph alive on the GPU and slowly eats VRAM until you hit out of memory. - Chasing a NaN before checking the learning rate: a rate that is too high is the cause nine times out of ten. Lower it first, then look at inputs and custom losses.
- Trusting a seed across machines: seeds repeat a run on the same setup, not across different GPUs, CUDA versions, or CPU versus GPU.
- Reaching for multi-GPU too early: get single-GPU training correct first. DDP wraps working code; it does not rescue broken code.
Best Practices
- Pick the device once at the top with the cuda / mps / cpu check, then move everything to that one variable. Never hardcode
"cuda". - Print the gradient norm every few steps during early training. It is the cheapest health dashboard you will ever build.
- Turn on mixed precision with
torch.autocastplustorch.amp.GradScalerfor a near-free speedup and memory saving on modern GPUs. - Keep an effective batch size with gradient accumulation when VRAM is tight, instead of just shrinking the batch and hurting stability.
- Seed everything and log versions (Python, NumPy, PyTorch, CUDA) so a run is reproducible on your own machine and diagnosable on someone else's.
Frequently Asked Questions
Do I need a GPU to follow this PyTorch GPU tutorial?
Not to learn it. Every debugging block here (device selection, shape errors, NaN loss, gradient norms, seeds) is tested on a plain CPU and runs on any laptop. You only need a GPU to see the actual speedups and to reproduce the CUDA out of memory message. Google Colab gives you a free GPU runtime, which is the easiest way to try the fast path without buying hardware.
What is the fastest fix for a CUDA out of memory error?
Lower the batch size first, since activation memory scales almost linearly with it. If a smaller batch hurts training stability, keep the same effective batch size using gradient accumulation: run several micro-batches, let their gradients pile up, and step the optimizer once. Mixed precision and gradient checkpointing free up more memory after that, and together these steps solve most PyTorch GPU memory problems without new hardware.
Why does my PyTorch loss become NaN?
The most common cause is a learning rate that is too high, which makes updates overshoot until the loss overflows to infinity and then NaN. Lower the learning rate first. Other causes are NaN or infinity in your input data, a log of zero or square root of a negative in a custom loss, and missing gradient clipping in models prone to spikes like RNNs.
Is mixed precision safe, or will it hurt accuracy?
For the vast majority of models it is safe and standard practice. Autocast keeps sensitive operations like the loss in 32-bit while running matrix multiplies in 16-bit, and GradScaler prevents tiny gradients from rounding to zero. Final accuracy is usually within noise of full precision while training runs faster and uses less memory. Watch for rare instabilities in custom numerical code.
Interview Questions on PyTorch GPU Training
These come from real screens and onsites. Practice answering before you read each answer.
Q: You get "Expected all tensors to be on the same device, but found cuda:0 and cpu." What is wrong and how do you fix it?
One tensor was left on the CPU while the model lives on the GPU, so PyTorch refuses to combine them. The message names both devices, which points you at the fix: find the stray CPU tensor and send it to the same device with .to(device). It usually hides in something created inside the loop, like class weights, a mask, or an index tensor. The habit that prevents it is choosing one device variable at the top and moving both model and every input batch to it.
Q: Your batch size no longer fits in VRAM but you need the larger effective batch for stable training. What do you do?
Use gradient accumulation. Split the large batch into micro-batches that do fit, call backward() on each so the gradients accumulate in the .grad buffers, and step the optimizer only after the last micro-batch. Scale each micro-batch loss by one over the number of accumulation steps so the summed gradient matches a single full-size batch. This gives you the statistics of a big batch with the memory footprint of a small one. Mixed precision and gradient checkpointing stack on top for more headroom.
Q: How would you diagnose a loss that suddenly becomes NaN a few hundred steps into training?
Work a checklist in order. Lower the learning rate first, since overshoot is the most common cause, and consider a warmup schedule. Then check the inputs for NaN or infinity that a divide-by-zero in preprocessing can introduce, and inspect any custom loss for a log(0) or a square root of a negative. Finally add clip_grad_norm_ to cap spikes, and print the gradient norm each step so you see the explosion coming before it turns into NaN.
Q: What is the difference between DDP and FSDP, and when would you choose each?
DistributedDataParallel keeps a full copy of the model on every GPU and splits the data, averaging gradients after each backward pass. Choose it when the model comfortably fits on one GPU and you just want to train faster across several. Fully Sharded Data Parallel shards the parameters, gradients, and optimizer state across GPUs, so choose it when the model is too large to fit on a single card. FSDP builds on the ZeRO ideas that DeepSpeed pioneered, and both use a fast collective library like NCCL underneath to move data between GPUs.
Q: You set a seed but a colleague on a different machine gets slightly different results. Why?
A seed makes a run repeatable on the same hardware and software, but it does not guarantee identical numbers across different GPUs, different CUDA or PyTorch versions, or CPU versus GPU, because the underlying math kernels differ and some GPU operations are nondeterministic by design. You can tighten this with torch.use_deterministic_algorithms(True), accepting a possible slowdown, and by seeding DataLoader workers. Even then, treat cross-machine reproducibility as approximate and always log your library and driver versions.
What’s Next?
You can now handle PyTorch GPU training with confidence: move both model and data to one device, read a CUDA out of memory error and fix it with a smaller batch or gradient accumulation, decode shape mismatches from the dimensions PyTorch prints, chase down a NaN loss starting with the learning rate, watch the gradient norm for exploding or vanishing signals, flip on mixed precision for a real speedup, and reason about DDP and FSDP when one GPU stops being enough. These are the exact skills that separate someone who has read about deep learning from someone who ships it.
Want the full roadmap from Python basics through deep learning and deployment? Head back to the Python + AI/ML tutorial series home to see where this post fits and what to read next.
Related Posts
- First Neural Network with PyTorch: tensors, autograd, and the training loop
- CPU vs GPU vs TPU: which hardware for which job
- PyTorch vs TensorFlow: choosing a framework
Further reading: PyTorch documentation is the authoritative source on this.
Related Posts
Previous: PyTorch Dataset and DataLoader: The Real Training Loop
Next: Python: PyTorch vs TensorFlow vs JAX, Deep Learning Frameworks Compared
Series Home: Python + AI/ML Tutorial Series

No comment