In this article
Before you start
This page is part of the PyTorch and machine learning (ML) coding round guide.
Why interviewers ask this
A broken training run costs hours before anyone notices, so a debugging task shows how a candidate works under uncertainty. Candidates report ML coding rounds where they read unfamiliar PyTorch code, find bugs in a prewritten data pipeline, or fix a model and its training process.
A weak answer changes the learning rate, adds layers or rewrites the loss at random. A strong answer follows a fixed order, says what each check would show if it failed, and stops at the first failure. A Staff-level answer also names the bugs that a model can compensate for, because, as Karpathy (2019) puts it, "Neural net training fails silently". Interviewers also read how you narrate; the code review interview questions cover explaining someone else's code aloud.
Debugging a training loop from first principles
- Training loop
- A training loop is the code that repeats four actions: take a batch of examples, run the model forward to get a loss, run backward to get gradients, and update the weights. A bug in any of the four can leave the loop running without any error message.
- Overfit-one-batch test
- The overfit-one-batch test is a sanity test that trains the model on a handful of examples until it memorises them. A correct model and loop drive that loss toward zero; if they cannot, there is a bug, and tuning the rest is wasted effort.
In plain words
Think of a factory line that makes a product from raw material. If the product comes out wrong, you do not adjust every machine at once. You check the raw material, then the first machine, then each machine in turn. The analogy breaks in one place: a broken machine in a training loop does not stop the line. It keeps producing output that looks plausible.
In one precise sentence, debugging a training run is shrinking the possible error surface with checks whose expected result you know in advance. Karpathy (2019) describes that surface as large, logical rather than syntactic, and very tricky to unit test, so each check below has a known right answer.
Step 1: Fix the seed and look at the data right before the model
Fix the random seed and turn off data augmentation first. Both remove sources of variation, so two runs of the same code give the same result and the pipeline has fewer places to hide a bug. Karpathy (2019) also says to visualise the data immediately before the line that calls the model, because that is the only place that shows what the network really receives.
Input: one batch, exactly as the model gets it. Output: your own judgement that the examples and labels match. Print a few examples with their labels, check shapes and value ranges, and look for flipped labels, a wrong scale or examples in the wrong order. A pipeline bug upstream of that line cannot be found by staring at the model. You now know the input is what you think it is.
Step 2: Check the loss at initialization
The starting loss of a classifier has a known value. For a softmax classifier with n classes and a correctly initialised final layer, the first loss should be near the negative log of one over n, as Karpathy (2019) puts it, which equals the natural logarithm of n. Ten classes give 2.3026, and a thousand classes give 6.9078.1
Input: the untrained model and one batch. Output: one number to compare with ln(n). If the first loss is far higher, the final layer starts with large outputs (a bad initial scale), the loss is summed over the batch instead of averaged (which gives the batch size times ln n), or the loss function receives the wrong input, such as probabilities instead of raw scores.
If it is far lower, the loss is usually averaged over the wrong dimension or computed on the wrong tensor. This check costs one forward pass and rules out a class of mistakes before any training happens. The labels themselves are checked in Steps 1 and 4, so this step tests only the final layer and the loss call. You now know they are wired as expected.
Step 3: Check that examples do not leak into each other
A dependency check is a test that takes the sum of the outputs of one example, runs backward all the way to the input, and confirms that only that example's input has a non-zero gradient; it charts which inputs an output depends on. The classic gradient check, which compares autograd gradients with finite differences (torch.autograd.gradcheck), is a different tool, used to test a custom backward pass. Karpathy (2019) describes this test for a bug where a reshape that should have been a transpose mixes information across the batch dimension, and the network still trains okay.
Input: a batch and an output of example i. Output: a gradient per input. In the toy model, the largest slope on the target input is 6.6 and exactly zero on the other three inputs; with batch mixing added, the other three inputs show 1.65 each.1 The same check shows whether an autoregressive model at time t depends only on earlier positions, which is the job of the mask in causal masks. For a transformer-specific version of this round, see the transformer debugging round. You now know the model treats each example on its own.
Step 4: Overfit one tiny batch
Take a few examples, as few as two, and train until the loss is near zero. Karpathy (2019) says to overfit a single batch of only a few examples; if the loss cannot reach about zero, there is a bug and you should not continue. In the toy run, four examples at a fixed step size of 0.5 fall from 1.0986 to 0.0055 in 200 steps, while the same loop with fresh random labels each step stays near 1.1
Input: a tiny batch, no augmentation, the same labels every step. Output: a loss curve that heads to zero. If it will not, suspect the wiring between the model and the loss, a step size that is too high or low, or labels that change. A common baseline setting is Adam at a learning rate of 3e-4, according to Karpathy (2019). Once this passes, move to the full data and regularise.
Step-by-step description
- Fix the seed and look at the batch.
- Compare with the log of n.
- Only example i gets a gradient.
- A tiny batch must reach near zero.
- Fix the seed and inspect the data right before the model.
- Compare the starting loss with the log of the number of classes.
- A gradient from one example's output must reach only that example's input.
- A correct loop drives the loss on a tiny batch toward zero.
Knowing it is half the job. Saying it is the other half.
Interview Clozo is an AI interview copilot for real interview rounds: it hears the question and puts a strong answer on your screen, and you say it in your own words. See the features or how it works.
Start Free Trial →Worked example
The numbers below come from a small script. The setup is illustrative: four examples, eight features, three classes, a linear softmax model, zero starting weights and a step size of 0.5.1 The starting loss is 1.0986, which is ln 3, and that matches the Step 2 rule.
With the correct labels the loss is 0.1141 after 10 steps, 0.0223 after 50 and 0.0055 after 200. With new random labels each step it measures 1.6818, 1.0137 and 1.4453 at those points, so it wanders near or above the starting value and never approaches zero.1 That gap is the signal: the first run passes the overfit-one-batch test and the second fails it.
| Check | Result (computed, illustrative setup) |
|---|---|
| Starting loss, 2 classes | 0.6931 |
| Starting loss, 10 classes | 2.3026 |
| Starting loss, 1,000 classes | 6.9078 |
| One-batch loss after 200 steps, correct labels | 0.0055 |
| One-batch loss after 200 steps, random labels | 1.4453 |
| Gradient on the other inputs, clean model | 0 on each |
| Gradient on the other inputs, batch-mixing model | 1.65 on each |
Text description of the diagram
- With correct labels the loss falls from 1.0986 to 0.0055 over 200 steps.
- With fresh random labels each step the loss stays near 1 and never approaches zero.
Memory is arithmetic too. For an illustrative model of 100,000,000 parameters stored in 32-bit numbers, the weights take 400,000,000 bytes and the gradients another 400,000,000. If the optimiser keeps two extra copies per parameter, as Adam does with its two moment estimates (Kingma and Ba, 2014, Algorithm 1), that adds 800,000,000 bytes, so the total is 1,600,000,000 bytes before any activations.1
Trade-offs and edge cases
Passing every check is not proof of correctness. Karpathy (2019) notes that a network can compensate for flipped labels, batch mixing and off-by-one leakage and still work fairly well, and that clipping the loss instead of the gradients silently ignores outlier examples. The checks reduce the error surface; they do not remove it. Also switch augmentation back on afterwards, because a bug that only appears with augmentation goes unnoticed while it is off.
A second trade-off is style. Write a specific, loopy version first, get it working, then vectorise it and confirm the results match, as Karpathy (2019) suggests. That costs time but gives you a reference. The same habit applies in the multi-head attention exercise.
The table lists the first check for each common symptom. A starting loss far from ln n points at the final layer scale or the loss call. A batch that will not overfit points at wiring, labels or step size. A model that overfits but stalls on the full run points at the data and augmentation. A run that runs out of memory early points at batch size or tensors holding history. A hang on many graphics processing units (GPUs) points at ranks that do not run the same collective calls.
Text description of the diagram
- A starting loss far from ln n points at the final layer scale or the loss call.
- A batch that will not overfit points at wiring, labels or step size.
- A hang on many GPUs points at ranks that do not run the same collective calls.
The bugs interviewers plant
A planted bug in a training loop is usually a one-line mistake that runs without an error and shows up only as a bad loss or a memory leak. Autograd records each operation as the forward pass runs so that the backward pass can replay it (Paszke et al., 2019), which is why several of these bugs involve the order of calls or a tensor that keeps its history. Reported rounds ask candidates to find bugs in a prewritten loop, so scan for these seven first.
| Bug | What goes wrong | Fix |
|---|---|---|
Missing optimizer.zero_grad() | Gradients accumulate across steps, so updates use stale sums | Clear gradients once per step, before loss.backward() |
optimizer.step() before loss.backward() | The update uses empty or stale gradients | Order: forward, loss, backward, step |
No model.train() / model.eval() switch | Dropout and batch norm behave wrongly in training or evaluation | Call the right mode before each phase |
Softmax before nn.CrossEntropyLoss | The loss expects unnormalized logits, so a softmax first distorts it (documentation) | Pass raw model outputs |
total_loss += loss | The sum keeps every step's graph alive and memory grows (PyTorch FAQ) | Add float(loss) or loss.item() |
Evaluation without torch.no_grad() | Graphs are built needlessly and memory use rises | Wrap evaluation in torch.no_grad() |
| Inputs and labels shuffled separately | Labels no longer match their examples | Shuffle indices once, or shuffle the pairs together |
How to answer it in an interview
The training loop runs but the model does not learn. How do you debug it?
I debug in a fixed order because training bugs fail silently. First I fix the seed, turn off augmentation and look at the batch right before the model. Then I compare the starting loss with the log of the number of classes. Then I take the gradient of one example's output back to the input and confirm only that example gets a gradient, which catches batch mixing. Finally I overfit a tiny batch; if the loss will not approach zero, I stop and look for a bug before tuning anything.
The spoken structure: say that bugs are silent, then name the four checks in order, then say what you do when one fails. For a longer build exercise, see the generative pre-trained transformer (GPT) from scratch walkthrough.
Common mistakes
- "The loss goes down, so the code is correct." (wrong) A network can train around a bug. (right) "A falling loss is necessary, not sufficient; I also run the leakage and overfit checks."
- "I'll tune the learning rate first." (wrong) That hides a pipeline bug behind a hyperparameter. (right) "I check the data and the starting loss before touching hyperparameters."
- "I debug with augmentation and several seeds on." (wrong) That adds variation and extra places for bugs. (right) "I fix the seed and turn augmentation off, then restore both."
- "I clip gradients by clipping the loss." (wrong) Clipping the loss ignores outlier examples. (right) "I clip the gradients after the backward pass and confirm which tensor I clipped."
Follow-up questions
How would you debug a PyTorch training loop running out of memory before the first epoch finishes?
Do the arithmetic first: parameters, gradients and optimiser state add up to a fixed amount. In the illustrative case above that is 1,600,000,000 bytes, and activations come on top.1 Then check what grows per step: a tensor still attached to its history keeps every intermediate value needed for backward alive. Log loss.item() or float(loss) instead of loss, as the PyTorch FAQ advises, run evaluation under undefined, and del large tensors you no longer need. Reduce the batch size to confirm the cause, and measure with torch.cuda.max_memory_allocated() instead of guessing.
How do you debug a training loop on a distributed GPU cluster?
First reproduce the bug on one device with a tiny batch, so you know the logic is right before it is spread across processes. Then check that every process runs the same sequence of collective calls, since one process that skips a call leaves the others waiting; PyTorch's distributed documentation describes the TORCH_DISTRIBUTED_DEBUG setting for extra checks on collective calls. For a sharded setup, see fully sharded data parallel (FSDP) explained.
What if the system is not returning model output?
Treat it as ordinary debugging of a call chain: check the status of the request, then timeouts, then how the response is parsed, then any asynchronous code that returns before the result arrives. Print what each stage receives and returns, and verify each stage in turn.
Why does the loss stop falling even though the code runs?
Run the overfit-one-batch test. If a tiny batch cannot reach near zero, suspect wiring, labels or step size; if it can, suspect the data, augmentation or regularisation on the full run.
How do you debug code you did not write?
Read it in the order the data flows, run the starting-loss check, and write a small focused function for the part you suspect, as Karpathy (2019) does before generalising. Narrate what you expect before each run.
- Training bugs are logical and silent, so a running loop proves little.
- Check in order: data and seed, starting loss, example leakage, one-batch overfit.
- A softmax classifier should start near ln(n); ten classes give 2.3026.
- A gradient from one example's output must reach only that example's input.
- If a tiny batch will not overfit, stop and find the bug before tuning.
Sources
- A Recipe for Training Neural Networks — Andrej Karpathy (blog), 25 Apr 2019
- Adam: A Method for Stochastic Optimization (arXiv 1412.6980) — arXiv (Kingma and Ba), 22 Dec 2014
- PyTorch: An Imperative Style, High-Performance Deep Learning Library (arXiv 1912.01703) — arXiv (Paszke et al.), 3 Dec 2019
- torch.nn.CrossEntropyLoss — PyTorch documentation, 1 Jan 2025
- PyTorch FAQ — PyTorch documentation, 1 Jan 2025
- torch.no_grad — PyTorch documentation, 1 Jan 2025
- Distributed communication package — PyTorch documentation, 1 Jan 2025
