Skip to content

The PyTorch training loop

  • Coding
  • Asked often
  • About 2 h 30 min to learn
  • 3 vetted resources
Build a plan that includes this

What this topic covers

The ML coding round, done in PyTorch: wrap data in a Dataset and DataLoader, define a model as an nn.Module, then write the loop that runs the forward pass, computes the loss, backpropagates and steps the optimiser, followed by an evaluation pass. Interviewers watch for the details that separate people who have trained models from people who have read about it: zeroing gradients, switching between train and eval modes, turning off gradient tracking for evaluation, and passing raw logits to the loss. The usual follow-up hands you a loss that won’t fall, and a strong answer debugs it methodically: check the loss at initialisation, overfit a single batch, inspect the data exactly as it enters the model, then look at the learning rate.

Builds on

You’re ready when you can

What interviewers check, from the resources below.

  • Write a custom Dataset with __len__ and __getitem__ and wrap it in a shuffled DataLoader
  • Write a training step in the right order: forward pass, loss, backward, optimiser step, zero the gradients
  • Write an evaluation loop that calls model.eval() and runs under torch.no_grad()
  • Explain why nn.CrossEntropyLoss takes raw logits rather than softmax outputs
  • Define an nn.Module with layers in __init__ and the computation in forward
  • Explain the difference between nn.BCEWithLogitsLoss and nn.BCELoss and why the logits version is more stable
  • Write a training function that moves each batch to the model’s device before the forward pass
  • Save and reload a model through its state_dict
  • State the loss you expect at initialisation for a C-class softmax classifier and why checking it catches bugs
  • Explain why overfitting a single batch is the first test of a new training loop

The PyTorch training loop interview questions

Answer out loud, then open the points a strong answer covers.

  • Walk me through a minimal PyTorch training and evaluation loop. What does each call do?
    • Each training step calls model.train(), runs the forward pass on a batch, computes the loss, calls loss.backward() for gradients, steps the optimiser, then zeroes the gradients before the next batch.
    • Gradients accumulate in each parameter’s .grad across backward calls, so forgetting optimizer.zero_grad() silently sums gradients from every previous batch into each update.
    • Evaluation calls model.eval() so dropout and batch norm switch to inference behaviour, and runs inside torch.no_grad() so no computation graph is built and memory stays low.
  • Your training loss isn’t falling at all. How would you debug it?
    • Check the loss at initialisation first: a C-class classifier with cross-entropy should start near log C, and a very different value points to a bug in the model, labels or loss.
    • Try to overfit a single small batch to near-zero loss; if the model cannot, the problem is in the code, such as a missing optimiser step, detached tensors or frozen parameters.
    • Inspect a batch exactly as it enters the model, confirm labels line up with inputs and that the loss gets logits not probabilities, then sweep the learning rate across orders of magnitude.

What to study, in order

Each link opened and checked by hand; the date says when.

  1. Start here

    Learn the Basics (opens in a new tab)

    PyTorch Tutorials · PyTorch · 60 min · Free · checked 26 Oct 2026

    Why
    The ML coding round expects you to write a Dataset, a DataLoader, a model and a train and test loop from memory, and this series is the canonical version of each.
    What to read
    From the series index, work through “Datasets and DataLoaders” (especially “Creating a Custom Dataset for your files” with __init__, __len__ and __getitem__, and “Preparing your data for training with DataLoaders”), “Build Model”, and “Optimization Loop” (“Loss Function”, “Optimizer” and “Full Implementation” with train_loop and test_loop). Then close the page and rewrite train_loop and test_loop on your own.
    What to skip
    “Tensors”, “Transforms” and “Automatic Differentiation” if you have done the backprop topic, and “Save, Load and Use Model”.

  2. Free alternative

    Tutorial 2: Introduction to PyTorch (opens in a new tab)

    UvA Deep Learning Tutorials · Phillip Lippe · 60 min · Free · checked 26 Oct 2026

    Why
    It builds one small classifier end to end, the same shape as a typical ML coding round, so you see every piece of the loop in context.
    What to read
    Read “Learning by example: Continuous XOR” in full: “The model” (nn.Module, Simple classifier), “The data” (The dataset class, The data loader class), “Optimization” (Loss modules, Stochastic Gradient Descent), “Training” including “Saving a model”, and “Evaluation”. Run the notebook if you can.
    What to skip
    “The Basics of PyTorch” section if you already know tensors and autograd, “Visualizing classification boundaries” and “TensorBoard logging”.

  3. Go deeper

    A Recipe for Training Neural Networks (opens in a new tab)

    Andrej Karpathy blog · Andrej Karpathy · 45 min · Free · checked 26 Oct 2026

    Why
    “My loss won’t go down, what do you do?” is the standard follow-up in the ML coding round, and this is the most cited checklist for answering it.
    What to read
    Read the two opening points (“Neural net training is a leaky abstraction” and “Neural net training fails silently”), then “The recipe” steps 1–3: “Become one with the data”, “Set up the end-to-end training/evaluation skeleton + get dumb baselines” (every bullet, such as verifying the loss at init and overfitting one batch) and “Overfit”.
    What to skip
    Steps 4–6 (“Regularize”, “Tune”, “Squeeze out the juice”) unless you have time; they matter less for a debugging question.

Companies that test Coding most

Put The PyTorch training loop in your plan

Tell us your interview date and hours a day; it lands on the right day, in the right order.

Build my plan