PyTorch
loss functions
model training
gradient computation
machine learning

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.

Practice ML system design

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:

python
1import torch
2import torch.nn as nn
3
4model = nn.Linear(3, 1)
5criterion = nn.MSELoss()
6
7x = torch.tensor([[1.0, 2.0, 3.0]])
8y = torch.tensor([[1.5]])
9
10prediction = model(x)
11loss = criterion(prediction, y)

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.

python
1import torch
2
3w = torch.tensor(2.0, requires_grad=True)
4x = torch.tensor(3.0)
5
6y = w * x
7loss = y ** 2
8
9print(loss)

Here, loss has a path back to w, so loss.backward() can compute dloss/dw.

python
loss.backward()
print(w.grad)

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.

python
for name, param in model.named_parameters():
    print(name, param.requires_grad)

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.
python
1a = torch.tensor(2.0, requires_grad=True)
2b = torch.tensor(5.0, requires_grad=True)
3
4loss = (a * 3) ** 2
5loss.backward()
6
7print(a.grad)
8print(b.grad)

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.

python
1optimizer = torch.optim.SGD(model.parameters(), lr=0.1)
2
3optimizer.zero_grad()
4loss.backward()
5optimizer.step()

So the workflow is:

  1. forward pass builds the graph,
  2. loss is computed from graph outputs,
  3. backward() computes gradients on participating tensors,
  4. 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.

python
prediction = model(x)
broken = prediction.detach()
loss = criterion(broken, y)

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.

Free course
Beginner
7 lessons
2 hours
Tackling System Design Interview Problems

A short course that equips you with the skills to approach system design interviews methodically.

Start the free course
Basis Lab
No 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 course
Track 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.

Practice ML system design