Skip to content

← Autograd Mechanics step 5 of 8

Medium Primitives

Backward twice

Jump to the task

Two losses can depend on the same forward computation. If you backpropagate through one and then the other, the second pass may need saved values that the first pass already released.

Keep the saved values for another pass

import torch

x = torch.tensor([1., 2.], requires_grad=True)
y = x.square()
y.sum().backward(retain_graph=True)
print(x.grad.tolist())  # [2.0, 4.0]
(y * 3).sum().backward()
print(x.grad.tolist())  # [8.0, 16.0]

The first loss contributes 2 * x. The second contributes 6 * x. Backward adds both contributions to x.grad, giving 8 * x because the initial gradient was None.

Without retention on the first pass, the second pass through the square operation raises because its saved values have been released. This does not mean all graph objects disappear. Some simple graphs with no saved values needed for backward can be traversed again without retention; do not rely on that for general computations.

Retaining a graph has a memory cost

By default, backward releases saved intermediates it no longer needs. Retaining them extends their lifetime and can increase memory use; the amount depends on the computation. Omitting retention on the final pass allows that pass to release its saved intermediates. Other references or unvisited graph branches may still keep data alive.

For the sum of these two losses, one backward call on their combined loss would normally suffice. This exercise asks for two calls specifically to practice retaining and reusing a shared graph.

For separate training steps, usually run a fresh forward pass for each step. For logging, accumulate detached loss values so the log does not keep references to computation graphs. Gradient accumulation over separate forward passes does not by itself require retaining each graph.

Higher derivatives need gradient recording

retain_graph=True preserves saved forward data. create_graph=True records operations used to compute a gradient so you can differentiate that gradient when it depends on the input:

x = torch.tensor([1., 2.], requires_grad=True)
loss = x.pow(3).sum()
g = torch.autograd.grad(loss, x, create_graph=True)[0]
second = torch.autograd.grad(g.sum(), x)[0]
print(second.tolist())  # [6.0, 12.0]

When retain_graph is left unspecified, its default follows create_graph. Retention alone does not create a differentiable gradient calculation.

Your task

def two_losses(x: torch.Tensor) -> torch.Tensor

The harness supplies a floating-point leaf x requiring gradients, with x.grad initially None. Compute y = x ** 2 once. Call backward first on y.sum() and then on (y * 3).sum(). Return the accumulated gradient without clearing it between passes. The expected result is 8 * x.