Table of Contents
What a Sanity Check Is and Why It Works
A sanity check is a deliberately small, controlled experiment that answers a simple question: can this exact model, with this exact training loop, reduce the loss on a tiny slice of data where success should be easy. When training fails, it is often unclear whether the problem is in the model, the data pipeline, the loss, the labels, the optimizer, or the loop itself. Overfitting a small batch removes most uncertainty by making the target behavior extremely strong and observable. If the model cannot memorize a handful of examples, something is wrong in the fundamentals of the setup, not in the choice of architecture or hyperparameters.
The Core Procedure: Overfit a Single Batch
Pick a very small number of examples, often one batch from your DataLoader such as 8, 16, or 32 items. Then train on that same batch repeatedly for many steps, without changing the data, until the model nearly memorizes it. In classification, memorization usually looks like training accuracy approaching 100 percent and loss approaching the theoretical minimum, which is close to 0 for common losses like cross entropy when predictions become confident and correct. In regression, memorization looks like training loss becoming extremely small and predictions matching targets closely on those few points.
During this test you want to eliminate sources of randomness and variation. Use a fixed batch, do not shuffle, and do not use data augmentation. If your pipeline normally includes augmentation, disable it temporarily because augmentation intentionally makes the input change each iteration, which directly fights memorization and can hide whether the training loop is correct.
Rule: If your model cannot overfit a single small batch, do not scale up training. Fix the pipeline, loss, labels, or loop first.
How to Set Up the Batch Correctly
A common failure mode is that the “single batch” is not actually constant. If you keep iterating through a DataLoader with shuffling enabled, you will see different samples. Instead, explicitly grab one batch once, store it, and reuse it. If your dataset returns a dictionary or multiple tensors, store the entire structure and move everything to the correct device once.
Make sure you are actually training. The model should be in training mode, and you should be updating parameters. It is easy to accidentally leave the model in evaluation mode from earlier checks, or to run the forward pass under a no grad context, which prevents learning.
Rule: For the overfit test, always call model.train() and ensure you are not using torch.no_grad() around the training step.
What to Turn Off (Temporarily)
Dropout and heavy regularization can make memorization slower. Weight decay can also resist memorization. For this specific test, you want the most direct learning signal possible. You do not have to remove these permanently, but consider disabling dropout, setting weight decay to 0, and keeping the model simple if you are diagnosing a failure.
Batch normalization deserves special attention. It behaves differently in train mode versus eval mode and depends on batch statistics. With very small batch sizes, its statistics can be noisy and can hinder fast memorization. If your network uses batch norm and the overfit test struggles, try increasing the batch size for the test or temporarily replace batch norm with a simpler alternative later. The key point is not to redesign the final model now but to isolate whether the training loop works.
What You Should Expect to See
You should see a consistent downward trend in training loss over repeated updates on the same batch. Some noise is normal, but the overall direction should be clear. If the loss plateaus quickly at a high value, or if it oscillates wildly, suspect a bug or instability.
In classification with cross entropy, the loss is $-\log(p_{y})$ averaged over the batch, where $p_{y}$ is the predicted probability of the true class. As the model memorizes, $p_{y} \to 1$ and the loss tends toward 0. If accuracy improves but loss does not, or loss decreases but accuracy stays random, suspect a mismatch between labels and loss, or a metric calculation bug.
Rule: On a fixed small batch with no augmentation, training loss should drop significantly, often by an order of magnitude, if the setup is correct.
If It Does Not Overfit: The Most Informative Failure Patterns
If the loss does not change at all, you are likely not updating parameters. Typical causes include missing optimizer.step(), forgetting loss.backward(), running under no grad, or using parameters that are not registered in the optimizer. Another possibility is that all gradients are zero due to a disconnected graph, which can happen if you accidentally detach tensors or convert them to Python numbers inside the forward pass.
If the loss becomes NaN or inf quickly, the learning rate may be too high, the inputs may be poorly scaled, or your loss receives invalid values. For example, passing raw logits into a loss that expects probabilities, or passing probabilities into a loss that expects logits, can cause numerical issues. Also check for NaNs in input tensors and targets.
If the loss decreases a bit but accuracy remains around chance for classification, labels may be wrong, shuffled relative to inputs, or in the wrong format. A very common mistake is using one hot labels with CrossEntropyLoss, which expects integer class indices, not one hot vectors. Another common mistake is an off by one label encoding, where classes are numbered differently than expected.
If accuracy improves but then collapses, you might still have nondeterminism from augmentation, shuffling, or a changing batch. Confirm that you are reusing exactly the same tensors each step.
Tightening the Test: Make It Even Easier
If the first attempt fails, make the problem easier so that success is unmistakable. Use an even smaller subset, such as 1 to 4 examples. Use a higher capacity model for the test, such as a slightly wider MLP for tabular data, while keeping everything else the same. Increase the number of training steps. Lower the learning rate if you see instability, or raise it slightly if learning is extremely slow and stable.
This is also the moment to simplify the input. If you have a complex preprocessing pipeline, test with raw tensors that you know are finite and correctly shaped. The goal is not to keep all production complexity, it is to prove the learning loop is capable of fitting known data.
Rule: When debugging, simplify until it works, then reintroduce complexity one piece at a time.
A Minimal Reference Training Step to Compare Against
When you run the overfit test, each iteration should conceptually follow the same structure: set train mode, compute predictions, compute loss, zero gradients, backpropagate, take an optimizer step. If your code differs substantially, that is fine, but you should be able to map it to this sequence. If you cannot, that is a sign that the loop may be doing something unintended.
Even if you do not change your implementation, it is useful to log a few values every N steps on the fixed batch, such as loss and a simple metric. If you also print the norm of a representative parameter tensor before and after training, you can confirm that parameters are changing, which is a fast way to catch a missing optimizer step.
Knowing When to Stop the Sanity Check
Stop when you have clear evidence of memorization or clear evidence of failure. If you can overfit a tiny batch, you have validated the end to end path from data to loss to gradients to parameter updates. You can then move to larger scale training with much higher confidence that issues like poor validation performance are about generalization, not about a broken training loop.
If you cannot overfit after simplifying and stabilizing, do not proceed to longer training runs. At that point, the most productive next step is to verify the exact mapping between inputs and targets, confirm that the loss matches the task, and confirm that gradients are flowing to the parameters you think you are training.