Understanding PyTorch models and loss functions

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

Why Predictions Are Not Training

Imagine a trained network that has already learned from many labeled examples. You now give it a new input X and ask for its output. The program runs the forward pass, which means the input flows through the layers of the model to produce a result, and then prints that result. The short snippet is simple, but two choices in it matter a great deal. The first is that the forward pass is wrapped in a special no_grad block. The second is that the printed values are treated as something other than final answers.

This article follows the program's execution step by step. First you will see what PyTorch does during an ordinary forward pass. Then you will see what changes when the same pass runs inside torch.no_grad(). Finally you will look at the output itself, which is a set of logits rather than a set of probabilities, and find the point where the programmer must convert one into the other.

Tracing the Forward Pass in Training Mode

During training, PyTorch records each operation applied to a tensor. That record is called the computational graph. The graph lets gradients flow backward through the same operations later, so that PyTorch can work out how each weight should change. This process is called backpropagation, and it is the reason the graph exists at all.

Read the comparison below as a trace of one forward pass in the ordinary mode. Every operation leaves a trace, and that trace is kept so that a backward pass can use it.

How does the same forward pass differ in what it records and keeps in memory when run with gradient tracking compared with inside a no_grad() block?

What torch.no_grad() Switches Off

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.

The key observation is about timing. 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 takes extra computation and additional memory, and neither cost buys the user anything. Inference is therefore handled differently from training, and no_grad() is the tool that does this.

torch.no_grad() 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 block is not a permanent change to the model. It only governs the operations performed inside it, so the same model can still be trained later.

Using no_grad as a Context Manager

The pattern below is a small illustration written for this article, not code from the source. It places a single forward pass inside the block and assigns its result to a variable. The model here is a tiny linear layer with three outputs, and the input is one example with four features.

python

import torch
import torch.nn as nn

model = nn.Linear(4, 3)   # small model with 3 outputs
x = torch.randn(1, 4)     # one input with 4 features

with torch.no_grad():
    logits = model(x)     # only the forward pass is inside the block

print(logits.requires_grad)

import torch
import torch.nn as nn

model = nn.Linear(4, 3)   # small model with 3 outputs
x = torch.randn(1, 4)     # one input with 4 features

with torch.no_grad():
    logits = model(x)     # only the forward pass is inside the block

print(logits.requires_grad)

Output (expected)

False

Notes on the code

import torch brings in PyTorch
nn.Linear(4, 3) builds a model whose parameters are tracked by default
with torch.no_grad(): starts the block where gradient tracking is off
logits = model(x) is the forward pass that produces the predictions
requires_grad reports whether the result is tracked for gradients

Notice the indentation. The line inside the block is the only line governed by no_grad(). Lines before and after the block run normally. This is why the source describes the block as wrapping the forward pass that produces predictions, and nothing else.

Tracing the Same Pass Without no_grad

To see the effect, remove the block and run the same forward pass in ordinary mode. The code is otherwise identical. Compare the two results side by side.

Before

model = nn.Linear(4, 3)
x = torch.randn(1, 4)

logits = model(x)

print(logits.requires_grad)

Forward pass is tracked, so the result is still connected to the graph and prints True

After

model = nn.Linear(4, 3)
x = torch.randn(1, 4)

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

print(logits.requires_grad)

Forward pass runs with tracking off, so the result is not tracked and prints False

What do you think happens?

Inside the with torch.no_grad() block, does logits.requires_grad print True or False?

Reveal answer

Answer:

False

Inside the block, PyTorch does not track gradients for the operations performed there, so the output is not marked as requiring gradients.

Logits Are Not Yet Probabilities

The output of the final layer is a set of raw scores called logits. 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.

PyTorch models commonly return logits, which are the raw outputs of the last layer without a nonlinear activation applied. Turning logits into class-membership probabilities for a multi-class problem requires the softmax function. For binary classification the analogous step is sigmoid.

Where does the softmax step sit between the model's raw output and the class-membership probabilities, and why do the logits not sum to 1?

In the pipeline, the logits are the value the model hands back. They are not yet probabilities, which is why they do not have to sum to one. The softmax call is the step that produces the probabilities, and it sits after the model rather than inside it.

Where the Softmax Call Belongs

PyTorch's usual loss functions apply softmax or sigmoid internally, combining that step with the negative log-likelihood loss in a single class. The source gives the reason as numerical efficiency and stability. Because of this, training code never needs to apply softmax explicitly. The loss function does it.

This arrangement creates a gap for inference. When you want to report how likely each class is for a given input, nothing in the model has performed that conversion. The source states plainly that you have to call the softmax function explicitly in that case. The flow below shows the order of steps during inference and where each piece applies.

In what order do the input, model layers, logits, softmax, and predictions flow through the code, and at which step does the no_grad() block apply?

Mistakes Beginners Make with no_grad and Logits

Common mistake
Assuming no_grad() changes the model permanently
Why it fails
The block only governs the operations performed inside it. Code outside the block behaves as usual, and the model can still be trained later.
Do this instead
Keep the with block around the inference forward pass and leave the model itself unchanged.
Common mistake
Wrapping training code in no_grad()
Why it fails
Training needs the computational graph so that backpropagation can compute how each weight should change. With tracking off, there is nothing to backpropagate through.
Do this instead
Use no_grad() only when the model is used for predictions after training.
Common mistake
Treating logits as final probabilities
What it looks like
Printing the raw output of the final layer and reporting it as the chance of each class
Why it fails
Logits can be negative and are not constrained to add up to one. The model stops before the probability conversion.
Do this instead
Call the softmax function explicitly on the logits when you need class-membership probabilities.
Common mistake
Adding an explicit softmax before a standard loss function during training
Why it fails
PyTorch's usual loss functions apply softmax or sigmoid internally, so the conversion would be repeated.
Do this instead
Pass the logits to the loss function and let it perform the softmax step.

Practice Tracing an Inference Block

MEDIUM

A colleague writes the following inference code. Identify which lines are covered by gradient tracking being turned off, and state where the softmax call should go if the goal is to report class probabilities. Explain why the training loss function would not need that call.

Hints

Only the indented lines under the with statement are inside the no_grad block.
The model returns logits, and the probability conversion is a separate step.
Think about whether the training loss function already applies softmax internally.

Your attempt (not saved)