How do loss functions know for which model to compute gradients in PyTorch?
ML System Design practice on Codemia
Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.
Introduction
In PyTorch, the loss function does not know about a model object in any special global sense. It only sees tensors. Gradients are computed for whatever tensors participated in building the loss and have requires_grad=True, which is why the real mechanism is the autograd graph, not a hidden pointer from the loss to “the current model.”
The Loss Operates On Tensors, Not On A Named Model
A typical training step looks model-centered:
It feels as if criterion knows which model is being trained. In reality, it does not. The loss was built from prediction, and prediction was built from operations involving the model parameters. That chain of tensor operations is what PyTorch tracks.
Autograd Builds A Graph During The Forward Pass
Each differentiable operation creates graph edges linking outputs back to inputs. The graph is dynamic and created during the forward pass.
Here, loss has a path back to w, so loss.backward() can compute dloss/dw.
No model class is required for this example. That is the clearest proof that autograd follows the graph, not a special loss-to-model relationship.
Model Parameters Are Just Leaf Tensors
When you define an nn.Module, its trainable parameters are tensors registered on the module.
During the forward pass, these parameters are used in operations. If the loss depends on the output of those operations, gradients flow back to the parameters.
So the real answer to the question is:
- the loss knows nothing magical about the model,
- the graph knows which tensors contributed to the loss,
- '
backward()traverses that graph and accumulates gradients in relevant leaf tensors.'
Why Only Some Tensors Get Gradients
A tensor gets gradients only if both of these are true:
- it was involved in producing the loss,
- it has
requires_grad=True.
a.grad is filled because a contributed to the loss. b.grad stays None because b never appeared in the computation.
That is why PyTorch does not need to ask which model to update. The graph already encodes which tensors matter.
The Optimizer Uses The Model Parameter List
There is a second layer that people often mix up with autograd: the optimizer. The optimizer does not discover parameters from the loss. You explicitly pass the parameters into it.
So the workflow is:
- forward pass builds the graph,
- loss is computed from graph outputs,
backward()computes gradients on participating tensors,- the optimizer updates only the parameter tensors it was given.
This separation is important. Autograd determines gradients. The optimizer determines which parameter list gets stepped.
Detached Tensors Break The Chain
If you detach a tensor, convert it to NumPy, or otherwise remove it from the graph, the loss can no longer propagate gradients through that path.
Calling loss.backward() here will not backpropagate into the model parameters through broken, because the connection was intentionally cut.
This is often the real reason gradients seem to “not know” about the model.
Common Pitfalls
- Thinking the loss function has a special reference to the model object.
- Forgetting that gradients only flow through tensors that contributed to the loss.
- Expecting gradients on tensors with
requires_grad=False. - Detaching tensors or converting them outside PyTorch before computing the loss.
- Confusing autograd’s graph traversal with the optimizer’s explicit parameter list.
Summary
- PyTorch loss functions operate on tensors, not on a hidden “current model.”
- Gradients are computed by traversing the autograd graph built during the forward pass.
- Model parameters receive gradients only if they participated in the loss computation and require gradients.
- The optimizer updates only the parameter tensors you passed into it.
- If the computation path is detached, gradients cannot flow back to the model.
Related reading
- How do you alter the size of a Pytorch Dataset?
- How do you invert a tensor of boolean values in Pytorch?
- How does one create a distributed data loader with PyTorch's TorchMeta for meta-learning?
- How does one implement adversarial examples in pytorch?
- How do recommendation systems work?
- How do sample weights work in classification models?
- How does one use Pytorch cuda with an A100 GPU?
- How does pytorch backprop through argmax?
.png&w=3840&q=75)
Tackling System Design Interview Problems
A short course that equips you with the skills to approach system design interviews methodically.
Start the free courseNo course covers this one yet
Basis Lab writes one for you from a sentence about what you want to be able to do, then teaches it and checks you understood. Your first course is free.
Build my courseTrack what you have practised
A free account saves your progress, solutions and study plan across every problem on Codemia.
ML System Design practice on Codemia
Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.