Table of Contents
Why shape and NaN checks come first
When a PyTorch model fails, the fastest wins usually come from two checks, tensor shapes and invalid numbers. Shape mistakes cause immediate runtime errors or subtle logic bugs where the code runs but learns nothing. NaNs and Infs often appear later, during training, and can silently corrupt the loss, gradients, and metrics. This chapter focuses on practical, repeatable ways to detect both early, before you spend time tuning optimizers, architectures, or data pipelines.
Verifying shapes systematically
The most reliable way to avoid shape confusion is to make shapes explicit at the boundaries between steps, dataset output, batching, model input, model output, and loss input. In PyTorch, the most common beginner mismatches happen at these boundaries, especially around batching and around classification targets.
A simple habit is to print or log the shape and dtype of tensors at the entry and exit of your model’s forward. For example, you want to know what x looks like when it arrives from the DataLoader, and what logits look like when they come out. When you do this, include the batch dimension and do not assume it is present. A dataset returns a single example, a dataloader returns a batch. If you test your model with one sample, you can accidentally feed a tensor shaped like a batchless example and later fail when training.
Rule: Decide and document one expected shape convention for each tensor, including the batch dimension, and check it at the model boundary. Most training bugs are mismatched conventions, not “bad learning rates”.
If you want a stricter approach, use assertions inside forward. Assertions are not elegant, but they catch mistakes at the exact point they occur, which is ideal for debugging. For example, if you expect a 2D tensor for an MLP, you can assert x.ndim == 2. If you expect image tensors, you can assert x.ndim == 4 and that the channel dimension is where you expect it. Avoid asserting fixed batch sizes, because the last batch can be smaller and evaluation often uses different batch sizes.
Also keep in mind that some layers impose specific shape constraints. nn.Linear expects the last dimension to match in_features, while earlier dimensions can be anything and are treated as batch like dimensions. This often surprises beginners because Linear can accept inputs shaped like (batch, features) and also (batch, time, features), but only if the last dimension is features. When debugging, checking x.shape[-1] right before the layer is often more informative than checking the full shape.
Typical shape contracts for common tasks
For regression with a single scalar target, it is common to produce outputs shaped (N, 1) or (N,). Either can work, but you must match the loss function expectations and keep it consistent across training and validation. Many PyTorch losses will broadcast silently, which can mask shape bugs, for example comparing (N, 1) predictions to (N,) targets might broadcast in ways you did not intend. A good debugging step is to print both shapes right before the loss and ensure they match exactly.
For multi class classification using nn.CrossEntropyLoss, the standard contract is that logits have shape (N, C) and targets have shape (N,) with integer class indices in the range [0, C-1]. Beginners frequently one hot encode targets and pass (N, C) targets into CrossEntropyLoss, which is not what it expects.
Important: nn.CrossEntropyLoss expects raw logits of shape $(N, C)$ and integer targets of shape $(N,)$. Do not apply softmax before CrossEntropyLoss.
For binary classification, you commonly choose between two patterns. Either use nn.BCEWithLogitsLoss with logits shaped (N,) or (N, 1) and targets as floats with the same shape, or treat it as two class classification with CrossEntropyLoss using logits (N, 2) and targets (N,). Both are valid, but mixing the conventions causes confusing shape and dtype errors. In particular, BCEWithLogitsLoss expects floating targets, not integer class indices.
Diagnosing shape issues with minimal disruption
When you hit a runtime error like “mat1 and mat2 shapes cannot be multiplied” or “Expected input batch size to match target batch size”, start by tracing the first tensor whose shape is not what you expected. Do not jump to rewriting layers. Instead, locate the exact point where the shape changes incorrectly. The easiest way is to insert temporary shape prints between key steps in forward, or use forward hooks to log layer input and output shapes. Hooks are useful when the model is large, but for beginners, a few well placed prints are often faster.
Another high value check is verifying that your DataLoader produces what you think it does. Print x.shape and y.shape for the first batch, and also print x.dtype, y.dtype, and y.min() and y.max() for classification. Many “model bugs” are actually dataset label bugs that manifest as shape problems or invalid values.
What NaNs and Infs are and how they appear
NaNs and Infs are usually symptoms, not root causes. Common sources include taking the logarithm of zero, dividing by zero, overflow from large exponentials, unstable losses, exploding activations, and overly large learning rates. They can appear in the forward pass, in the loss, or in gradients during backpropagation. Once a NaN enters parameters or optimizer state, it often spreads quickly.
PyTorch provides direct checks for invalid values. torch.isnan(t) identifies NaNs, torch.isinf(t) identifies infinities, and torch.isfinite(t) checks for both. A practical pattern is to check torch.isfinite(loss) every iteration during early debugging, and to stop immediately when it fails, while you still have context.
Rule: The first moment you see a NaN or Inf, stop and inspect. Continuing training usually destroys the evidence of where it started.
Where to check for NaNs: inputs, outputs, loss, gradients
A disciplined approach checks in four places. First check model inputs from the dataloader, because data preprocessing can introduce NaNs, especially when normalizing by a standard deviation that is zero or when parsing missing values. Second check model outputs, because invalid operations in forward can introduce NaNs. Third check the loss, because loss functions can blow up when given invalid targets or extreme logits. Fourth check gradients, because exploding gradients can appear even if the forward pass is finite.
In practice, you can start with checking loss finiteness each step. If it fails, then check outputs and inputs. If loss is finite but parameters become NaN after optimizer.step(), then check gradients and also consider optimizer settings. Some optimizers can amplify issues if the learning rate is too high.
Using anomaly detection to pinpoint the operation
PyTorch can help you locate the exact operation that created a NaN in the backward pass using anomaly detection. Enabling it makes training slower, but during debugging it is valuable because it produces a stack trace pointing to the problematic operation. Use it temporarily, reproduce the NaN quickly, then turn it off.
An important nuance is that anomaly detection is primarily for backward pass issues. If NaNs appear in the forward pass, you often need forward checks, such as verifying torch.isfinite on intermediate tensors or using hooks to inspect layer outputs.
Common NaN patterns in beginner models
One frequent pattern is applying log to a probability that can be zero. If you compute something like torch.log(p) without ensuring $p > 0$, NaNs are expected. Use numerically stable formulations, or clamp inputs with a small epsilon when appropriate. Another pattern is applying softmax and then log manually for classification instead of using CrossEntropyLoss, which already combines log_softmax and negative log likelihood in a stable way.
Another pattern is normalization. If you normalize with (x - mean) / std and std is zero for some feature, the result is Inf or NaN. This can happen with constant columns in tabular data or images with uniform regions if you compute per sample statistics incorrectly.
A third pattern is learning rate too high. The loss may be finite for a few steps and then suddenly become NaN. In that case, the first checks are to reduce the learning rate and to inspect gradient norms, which is covered in later chapters, but the key debugging step here is simply confirming whether NaNs first appear in gradients or in the forward outputs.
A minimal “stop on NaN” guardrail
During early training runs, add a guardrail that stops training as soon as the loss becomes non finite. The goal is not elegance, it is fast feedback. You can also save the batch that caused the failure, by keeping a reference to the current x and y and writing them to disk, so you can reproduce the issue without rerunning the entire epoch.
Important: A training loop that ignores NaNs can keep running and produce meaningless checkpoints. Always validate that loss and key metrics are finite.
Interpreting what you find
If shapes are wrong, fix the contract at the boundary, usually in one of three places: dataset output, collate and batching, or the first and last layers of the model. Avoid scattering view and squeeze calls throughout the model to “make it work”, because that often creates silent bugs when batch size changes or when a dimension equals 1.
If NaNs appear in inputs, the issue is in data loading or preprocessing. If NaNs appear in model outputs but not inputs, the issue is in forward, often an unstable operation. If NaNs appear only after backprop or optimizer step, the issue is likely gradients or optimizer dynamics, which you will handle with gradient checks, clipping, and learning rate adjustments later. The key outcome of this chapter is that you can locate which stage first becomes invalid, and you can do it quickly and repeatedly.