The Deep Learning Mental Model
PyTorch's tensors, its autograd engine, and the training loop can look like three separate topics learned one after another.
//
Search across all documentation pages
PyTorch's tensors, its autograd engine, and the training loop can look like three separate topics learned one after another.
They are actually one machine: a computation graph built from tensor operations, a backward pass that walks that graph to compute gradients, and a training loop that repeatedly uses those gradients to nudge parameters toward lower loss.
This page connects PyTorch Basics, Autograd, Training Loops, and GPUs & Mixed Precision into one model, rather than four separate topics to memorize independently.
A tensor is PyTorch's core data structure, an n-dimensional array much like a NumPy array, but with two additions that matter for deep learning: it can live on a GPU, and it can track the operations applied to it.
That second property is what makes autograd possible: every operation performed on a tensor with requires_grad=True is recorded into a computation graph, a directed graph where nodes are operations and edges are the tensors flowing between them.
A useful analogy is a recipe with a receipt attached.
The forward pass is following the recipe, mixing tensors through operations to produce an output, and the receipt is the computation graph, listing exactly which operations happened in which order so you can later walk backward through it to see how much each ingredient contributed to the final result.
The forward pass is ordinary computation: input tensors flow through layers such as matrix multiplies, activations, and normalization, to produce an output like a predicted class or a next-token probability.
The backward pass is where autograd's real function shows up.
Calling .backward() on a scalar loss walks the recorded graph in reverse, applying the chain rule at each node to compute the gradient of the loss with respect to every parameter that contributed to it.
The training loop is this pair of passes repeated many times, and every step in it exists to answer a specific question in a fixed order:
for each batch:
forward pass -> what did the model predict? (build graph)
compute loss -> how wrong was the prediction?
backward pass -> how much did each parameter contribute to that error? (walk graph)
optimizer step -> nudge each parameter opposite its gradient
zero gradients -> clear old gradients before the next batch
The optimizer step - SGD, Adam, and their variants - is where the actual learning happens, taking the gradients autograd computed and updating each parameter by some function of that gradient, typically scaled by a learning rate.
Zeroing gradients exists because PyTorch accumulates gradients into .grad by default rather than overwriting them.
That design supports techniques like gradient accumulation across multiple small batches, but it silently corrupts a normal training loop if you forget to clear gradients before the next backward call.
for batch in loader:
optimizer.zero_grad() # clear last step's gradients
output = model(batch.x) # forward pass builds the graph
loss = loss_fn(output, batch.y)
loss.backward() # walk the graph backward
optimizer.step() # update parameters using .gradThe computation graph is rebuilt from scratch on every forward pass and discarded after .backward() by default, which is why PyTorch is described as define-by-run.
The graph is a record of what actually happened on this specific forward pass, not a static structure declared up front, which is what makes data-dependent control flow straightforward compared to frameworks requiring a graph defined before any data flows through it.
GPUs matter here because the forward and backward passes are almost entirely matrix multiplications and element-wise operations, which parallelize extremely well across a GPU's thousands of cores compared to a CPU's much smaller number of cores optimized for sequential work.
Mixed precision training exploits the fact that most of this arithmetic doesn't need full 32-bit floating point accuracy to converge.
Running matrix multiplies in 16-bit, or bfloat16, while keeping a 32-bit master copy of the parameters cuts memory use and often roughly doubles throughput on modern GPU tensor cores, with a gradient scaler used to prevent small gradient values from underflowing to zero in the lower-precision format.
At scale, the computation graph's memory cost becomes the binding constraint rather than raw compute, since every intermediate activation from the forward pass has to stay in memory until the backward pass consumes it.
Deeper models and larger batch sizes cost memory roughly linearly in both dimensions as a result.
Gradient checkpointing trades compute for memory directly: instead of storing every intermediate activation, it stores only a subset and recomputes the rest during the backward pass, a deliberate application of the same forward/backward mental model rather than a separate technique.
Transfer learning and fine-tuning reuse this model too.
Freezing a pretrained network's early layers with requires_grad=False removes those parameters from the graph's backward walk entirely, so gradients are only computed - and only optimizer steps applied - for the layers actually being adapted, which is what makes fine-tuning cheaper than training from scratch.
Distributed training extends the same loop across multiple GPUs or machines: each device runs its own forward and backward pass on a slice of the batch, and the gradients, not the data or the model, are what gets synchronized between devices before each optimizer step.
This is why network bandwidth for gradient synchronization, not raw compute, is often the bottleneck at scale.
| Precision / Scale Strategy | Strength | Weakness | Best Fit |
|---|---|---|---|
| Full 32-bit precision, single GPU | Simplest to reason about; no numerical surprises | Slowest option; caps model/batch size to one GPU's memory | Small models, debugging, establishing a correctness baseline |
| Mixed precision, single GPU | Meaningful speedup and memory savings on modern tensor cores | Adds a gradient scaler and occasional numerical instability to reason about | Most production training once correctness is established |
| Distributed (multi-GPU) training | Scales model and batch size beyond one device's memory | Gradient synchronization becomes a real bottleneck; adds operational complexity | Models or datasets too large for one GPU regardless of precision |
.backward() twice in a row just recomputes the same gradients." Without clearing them first, gradients accumulate into .grad rather than being overwritten, so a missing zero_grad() silently corrupts a standard training loop.PyTorch walks the computation graph recorded during the forward pass, in reverse, applying the chain rule at each operation to compute the gradient of the loss with respect to every tensor that has requires_grad=True and contributed to that loss.
PyTorch accumulates new gradients into a tensor's existing .grad attribute by default rather than overwriting it. Without clearing it first, gradients from the previous batch would still be added into the current batch's gradients, corrupting the update.
A directed graph recorded during the forward pass, where each node is an operation applied to a tensor and each edge is a tensor flowing between operations. It's what autograd walks backward to compute gradients, and it's rebuilt from scratch on every forward pass.
Because the computation graph is built dynamically as operations actually execute on real data, rather than declared statically before any data flows through it. This makes data-dependent control flow, like a different number of loop iterations per input, straightforward to express.
It updates each parameter using some function of its gradient and a learning rate. SGD subtracts a scaled gradient directly, while adaptive optimizers like Adam also track per-parameter running statistics of past gradients to adjust the effective step size.
Forward and backward passes are dominated by matrix multiplications and element-wise operations, which parallelize extremely well across a GPU's large number of simple cores, unlike more sequential workloads that suit a CPU's smaller number of powerful cores better.
It reduces memory use and increases throughput by running most arithmetic in a lower-precision format, 16-bit or bfloat16, instead of full 32-bit floats, which modern GPU tensor cores execute significantly faster, while keeping a 32-bit master copy of parameters for numerical stability.
Small gradient values can underflow to zero when represented in 16-bit floating point, silently losing information. A gradient scaler multiplies the loss by a scaling factor before the backward pass, and unscales gradients before the optimizer step, to keep values in a representable range.
Setting requires_grad=False on a layer's parameters removes it from the backward pass's gradient computation entirely. Autograd never computes a gradient for it and the optimizer never updates it, so both the backward pass and the update step are cheaper.
Each device computes its own forward and backward pass on its slice of the batch, and it's the resulting gradients, not the raw data or the model weights mid-training, that get synchronized, typically averaged, across devices before each optimizer step.
It stores only a subset of forward-pass activations instead of all of them, recomputing the rest during the backward pass when needed. It's useful when a model is memory-bound rather than compute-bound, since GPU memory is often the harder limit to work around.
Not unconditionally. Larger batches better utilize GPU parallelism up to a point, but they also increase the memory needed to store activations for the backward pass, and very large batches can change training dynamics rather than simply training faster for free.
Stack versions: This page was written for PyTorch 2.6+ and Python 3.14 (stable) / 3.13 (maintenance).
Reviewed by Chris St. John·Last updated Jul 15, 2026