Neural networks, deep learning, NLP · insight-0401
Assert prediction and target shapes right before every loss call
- Evidence type
- documented fact
- Difficulty
- easy
- Roles
- Data scientist, ML engineer
- Industries
- aviation
Summary: situation, action, reason, boundary
Situation
A regression network returns predictions of shape (N, 1) while the target tensor is (N,), for example predicted load factor per flight-day.
Action
Put
assert pred.shape == target.shape(or an explicitsqueeze(-1)) immediately before each loss and metric, and log both shapes once at step one.Reason
PyTorch follows NumPy broadcasting: trailing dimensions that are equal, 1 or missing expand silently, so a compatible-but-wrong shape never raises. The built-in MSE loss only warns; a hand-written loss gives no signal.
Boundary
Any elementwise loss or metric in PyTorch, NumPy or JAX, including evaluation code.
Symptom or risk
Broadcasting expands the difference to (N, N), so the loss compares every prediction with every target. Training runs, the loss falls a little, and the model converges towards the mean.
Exceptions
Losses with documented different shapes, such as cross-entropy taking (N, C) logits with (N,) class indices.
Trade-offs
One line per loss call; it must survive refactors, so keep it in the training step rather than a notebook cell.
Pitfalls
Flattening everything with view(-1) also flattens multi-output heads and hides a different bug.
Example
Synthetic illustration: with batch 256 the flawed loss averages 65,536 pairwise errors and validation MAE stalls at roughly the error of predicting the mean load factor.
Lessons and concepts
Claims and their sources
PyTorch broadcasts tensors whose sizes, compared from the trailing dimension, are equal, 1, or absent, expanding them without copying data.