# The PyTorch training loop

> What interviewers check on The PyTorch training loop, and what to study, in order.

Canonical: https://saadho.arushk.tech/topics/pytorch-training-loop
Last updated: 2026-10-26

Coding. Asked often. About 2 h 30 min to learn. 3 vetted resources.

## 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

- [Deep learning basics](https://saadho.arushk.tech/topics/deep-learning-basics)
- [Python for ML interviews](https://saadho.arushk.tech/topics/python-for-interviews)

## You’re ready when you can

- 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

**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

### 1. Learn the Basics

Start here. PyTorch Tutorials, PyTorch. Guide, about 60 min. Free. Last checked 2026-10-26. [Open](https://docs.pytorch.org/tutorials/beginner/basics/intro.html)

- **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.
- **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.
- **Skip:** “Tensors”, “Transforms” and “Automatic Differentiation” if you have done the backprop topic, and “Save, Load and Use Model”.

### 2. Tutorial 2: Introduction to PyTorch

Free alternative. UvA Deep Learning Tutorials, Phillip Lippe. Guide, about 60 min. Free. Last checked 2026-10-26. [Open](https://uvadlc-notebooks.readthedocs.io/en/latest/tutorial_notebooks/tutorial2/Introduction_to_PyTorch.html)

- **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.
- **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.
- **Skip:** “The Basics of PyTorch” section if you already know tensors and autograd, “Visualizing classification boundaries” and “TensorBoard logging”.

### 3. A Recipe for Training Neural Networks

Go deeper. Andrej Karpathy blog, Andrej Karpathy. Article, about 45 min. Free. Last checked 2026-10-26. [Open](https://karpathy.github.io/2019/04/25/recipe/)

- **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.
- **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”.
- **Skip:** Steps 4–6 (“Regularize”, “Tune”, “Squeeze out the juice”) unless you have time; they matter less for a debugging question.

## Companies that test this area most

- [Juspay](https://saadho.arushk.tech/companies/juspay)
- [Databricks](https://saadho.arushk.tech/companies/databricks)
- [Stripe](https://saadho.arushk.tech/companies/stripe)
- [ByteDance](https://saadho.arushk.tech/companies/bytedance)
- [Anysphere (Cursor)](https://saadho.arushk.tech/companies/anysphere)
- [Samsung Electronics](https://saadho.arushk.tech/companies/samsung)

## Related

- [All topics](https://saadho.arushk.tech/library)
- [Build a plan that includes this](https://saadho.arushk.tech/start)