PyTorch: A Flexible Deep Learning Framework

PyTorch builds its computation graph as your code runs rather than requiring you to declare it upfront. Ordinary Python control flow works, a standard debugger works, and printing a tensor mid-forward-pass shows you a real value. That is most of why it became the default for research, and the production gap it once had has largely closed.

The training step, line by line

Five lines carry the whole thing, and the order is not arbitrary:

for batch, targets in loader:
    optimizer.zero_grad()            # clear gradients from the previous step
    outputs = model(batch)           # forward pass; autograd records the graph
    loss = criterion(outputs, targets)
    loss.backward()                  # walk the graph backwards, accumulate .grad
    optimizer.step()                 # apply the update using those gradients

zero_grad comes first because PyTorch accumulates gradients rather than replacing them. That is a deliberate design choice — it is what lets you simulate a large batch on small hardware by running several forward and backward passes before stepping — but it means forgetting the call does not raise an error. It silently trains on the sum of every gradient since the start, and the model quietly fails to converge.

The forward pass builds the graph. Every operation on a tensor with requires_grad=True records itself and its inputs, so loss.backward() has a path to walk backwards, applying the chain rule to compute each parameter's gradient. Nothing is written to weights during backward; it only populates .grad.

optimizer.step() reads those gradients and updates the parameters. The optimiser must have been constructed over the same parameters the loss depends on — a separate module created after the optimiser will not be trained, silently.

Autograd: what is tracked and when to stop tracking

The graph exists to compute gradients, and holding it costs memory. Two places you must switch it off:

model.eval()                    # switches dropout and batchnorm to inference behaviour
with torch.no_grad():           # stops graph construction, frees a large amount of memory
    for batch, targets in val_loader:
        preds = model(batch)

Those two lines do different jobs and both are needed. eval() changes layer behaviour — dropout stops dropping, batch norm uses running statistics instead of the batch's. no_grad() stops the graph being built. Omitting eval() gives you validation numbers that are wrong but plausible, which is the worst kind of bug. Omitting no_grad() merely wastes memory until you run out.

The mirror-image mistake is accumulating the loss tensor itself into a running total. total += loss keeps the entire graph for every batch alive; total += loss.item() keeps a float. This is the most common cause of memory that grows through an epoch.

Eager, compiled, exported

Three modes, three purposes.

Eager is the default: operations run as Python reaches them. Maximum debuggability, and per-operation overhead. This is where you should develop.

Compiled — model = torch.compile(model) — traces the model and generates fused kernels. Typically a solid speedup on modern GPUs for one line of code, with a compile cost on the first batch and re-compilation whenever input shapes change. Worth turning on for training runs long enough to amortise it.

Exported — TorchScript, ONNX, or the newer export path — produces an artifact that runs without Python, which is what a C++ service, a mobile target, or a non-Python runtime needs.

Develop eager, compile when the run is long, export when the deployment target demands it. Exporting early costs you debuggability for a benefit you may not need.

Failures that do not raise

Deep learning bugs mostly do not crash; they produce a model that trains to a mediocre number and gives no clue why. Four recur:

Broadcasting. Shapes (N, 1) and (N,) broadcast to (N, N) without complaint. A loss computed over that is meaningless and looks fine. Assert shapes at boundaries.

Device mismatch. This one does raise, and the fix is to keep the device in one variable and move both model and data with it rather than scattering .cuda() calls.

Not shuffling the training loader. If the data is ordered by class, the model sees long runs of one label.

Learning rate. More training failures come from this single number than from architecture. If loss goes to NaN it is too high; if loss barely moves it is too low. Check it before you change the model.

When something is wrong, overfit a single batch first. A model that cannot drive the loss to near zero on ten examples has a bug in the code, not a shortage of data — and that test takes thirty seconds instead of a training run.

Common mistakes

Forgetting zero_grad. No error, no convergence.

Validating without eval(). Wrong numbers that look reasonable.

Accumulating loss tensors instead of floats. Memory grows until the epoch dies.

Tuning architecture before the learning rate. The learning rate explains more failures than everything else combined.