Backpropagation and the computational graph

torch.no_grad() is a context manager that tells PyTorch not to track gradients, which saves memory and computation when a model is used only for inference.

Programming

A Prediction Script With Two Choices

Imagine a short script that loads a trained network, passes it an input X, and prints the output. The script looks almost trivial, and that is the point. Nearly everything interesting in it is a decision made by the programmer rather than something the framework forces on you. There are two decisions to watch for. The first is wrapping the forward pass in a block that turns gradient tracking off. The second is treating the printed values as something other than final answers.

This article follows those two decisions. It starts with what PyTorch records during an ordinary forward pass, then shows how torch.no_grad() changes that recording, and finally explains why the printed numbers are logits that still need a separate conversion step before they can be read as probabilities.

What the Computational Graph Records

During training, PyTorch records each operation applied to a tensor so that gradients can later flow backward through the same operations. That recorded structure is the computational graph. It connects the input to the output through the sequence of operations that produced the output, and it is what makes backpropagation possible. Backpropagation uses this structure to compute how each weight should change.

When gradient tracking is on, which nodes and links does PyTorch record as the input moves to the output?

Read the diagram as a trace of the forward pass. Each operation is a link in the chain, and each link points from an input to the result it produced. When tracking is on, PyTorch keeps this chain so that a later backward pass can walk it in reverse. The chain is the reason gradients can reach the weights that produced the output.

Why Inference Does Not Need the Graph

Gradient tracking exists to support backpropagation during training. When a model is used only to make predictions after training, nothing will ever call backpropagation on that forward pass. Recording the graph in that situation is wasted work. It costs extra computation and additional memory, and neither cost buys the user anything. This is the reason inference should be handled differently from training.

Definition

torch.no_grad() is a context manager that tells PyTorch not to track gradients, which saves memory and computation when a model is used only for inference.

Turning Off Tracking With no_grad

torch.no_grad() is a context manager, so it is used with a with statement. Everything inside the indented block runs with gradient tracking turned off. Code outside the block behaves as usual. The pattern places a single forward pass inside the block and assigns its result to a variable.

Code walkthrough

import torch

# model is a trained network and X is an input tensor
with torch.no_grad():
    logits = model(X)

print(logits)

1
The import line makes PyTorch available.
2
The with statement turns gradient tracking off for every operation in its indented body.
3
The forward pass model(X) runs inside the block, so no computational graph is recorded for it.
4
The printed value is the output of the last layer, which is the logits tensor.

with torch.no_grad(): opens a block in which gradient tracking is off
logits = model(X) runs the forward pass inside the block and stores its result
print(logits) sits outside the block and only reads the stored result

In the left column, every operation leaves a trace that a backward pass could later use. In the right column, the operations still compute the same numbers, but the trace is not kept. The output of the forward pass is identical in both cases; what differs is the bookkeeping that would have supported backpropagation.

Worked Example: Predicting Without Tracking

Worked example

Producing a prediction block that does not track gradients

A trained model is applied to a batch of inputs X to produce predictions for reporting. Write the forward pass so that PyTorch does not record gradient information for it.

1
Identify the forward pass
The only operation that produces predictions is model(X). Nothing will call backpropagation on it, so it is the line to wrap.
2
Open the no_grad block
Place the forward pass inside a with torch.no_grad(): block so that every operation in its body runs with gradient tracking off.
3
Assign the result
Store the output in a variable, such as logits, so the rest of the program can use it after the block ends.
4
Keep the rest outside
Code that does not compute the prediction, such as printing or saving results, can stay outside the block and behave as usual.

Result

The prediction is computed without a recorded graph, and the stored logits can be printed or passed on to later steps.

Logits Are Not Yet Probabilities

PyTorch models commonly return logits, which are the raw outputs of the last layer without a nonlinear activation applied. The output of the final layer is a set of raw scores. Those scores can be negative, and they are not constrained to add up to one. Many people expect a classifier's output to already be a probability distribution, but in the usual PyTorch style the model stops before that conversion. The conversion is a separate step that the programmer chooses to perform.

The pipeline shows the two routes. During training, the logits go to a loss function, and that loss function applies the softmax step itself. For reporting how likely each class is, the logits must pass through a softmax call that the programmer writes explicitly. Nothing in the model performs that conversion for you.

Where Softmax Must Be Called Explicitly

PyTorch's usual loss functions apply softmax or sigmoid internally, for numerical efficiency and stability. For binary classification the analogous step is sigmoid. These loss functions combine the softmax or sigmoid operation with the negative log-likelihood loss in a single class. This arrangement creates a gap for inference: training code never needs to apply softmax explicitly, but reporting class likelihoods does.

python

import torch

with torch.no_grad():
    logits = model(X)

# the model returned raw scores; convert them explicitly
probs = torch.softmax(logits, dim=1)
print(probs)

import torch

with torch.no_grad():
    logits = model(X)

# the model returned raw scores; convert them explicitly
probs = torch.softmax(logits, dim=1)
print(probs)

Notes on the code

logits is the raw output of the model, produced inside the no_grad block
torch.softmax(logits, dim=1) is the explicit conversion step the programmer must call
dim=1 names the axis over which the class scores are normalized

What do you think happens?

Before running softmax, do the raw logits for a classifier already add up to one?

Yes, the model outputs probabilities
No, the logits are raw scores and are not constrained to sum to one

Reveal answer

Answer:

No, the logits are raw scores and are not constrained to sum to one

The final layer produces raw scores that can be negative. The conversion to probabilities is a separate step that the programmer performs.

Scope of the no_grad Block

The block wraps only the forward pass that produces predictions. Lines outside the block keep their usual behavior. The no_grad block is not a permanent change to the model; it governs only the operations performed inside it. If the same model is needed for training later, nothing has to be undone.

python

import torch

x = torch.tensor([1.0, 2.0], requires_grad=True)

with torch.no_grad():
    y = x * 3    # computed inside the block

z = x * 3        # computed outside the block

import torch

x = torch.tensor([1.0, 2.0], requires_grad=True)

with torch.no_grad():
    y = x * 3    # computed inside the block

z = x * 3        # computed outside the block

Notes on the code

x is created with gradient tracking enabled
y is computed inside the block, so PyTorch does not track it
z is computed outside the block, so tracking behaves as usual

Mistakes Beginners Make with Inference Code

Common mistake
Assuming the printed output is already a probability distribution
What it looks like
Printing logits and reading the values as class probabilities.
Why it fails
The final layer produces raw scores that can be negative and are not constrained to add up to one.
Do this instead
Apply torch.softmax explicitly when you need class-membership probabilities.
Common mistake
Calling softmax before passing outputs to a training loss
What it looks like
Converting logits to probabilities and then giving those probabilities to a loss function that already applies softmax internally.
Why it fails
The usual loss functions apply softmax or sigmoid themselves, so the conversion is performed twice.
Do this instead
Pass the logits to the loss function and let it perform the softmax step.
Common mistake
Wrapping the entire training loop in no_grad
What it looks like
Placing the backward pass inside a no_grad block.
Why it fails
Gradient tracking exists to support backpropagation during training, so turning it off during training removes what backpropagation needs.
Do this instead
Use no_grad only around forward passes that produce predictions for inference.
Common mistake
Forgetting that code outside the with block still tracks gradients, so the block must be placed around the prediction line itself

Practice: Reading Tracking Scope

MEDIUM

A program computes predictions with a trained model inside a with torch.no_grad(): block, then applies torch.softmax to the logits and prints the result. Later in the same program, a training step runs a new forward pass outside any no_grad block and calls a loss function. Which forward pass records a computational graph, and which code performs the explicit probability conversion?

Hints

Check which lines are indented under the with statement.
Recall that the usual loss functions apply softmax internally.
Ask whether the explicit softmax call happens inside or outside the block, and whether that matters for the conversion.

Your attempt (not saved)

Key Takeaways

The computational graph is the record of operations PyTorch keeps during a tracked forward pass, and backpropagation uses it to compute how weights should change.
torch.no_grad() is a context manager that turns gradient tracking off for operations inside its with block, saving memory and computation during inference.
Models commonly return logits, which are raw scores that are not constrained to sum to one; the usual loss functions apply softmax or sigmoid internally.
When class-membership probabilities are needed for reporting, the programmer must call softmax explicitly on the logits.

Related concepts

Loss functions for classification

More concepts starting with B