Skip to content

← Memory, Copies and In-Place step 5 of 6

Medium Primitives

Updating a parameter by hand

Jump to the task

A trainable parameter is usually a leaf tensor: it requires gradients but has no operation recorded as its creator. With gradient recording enabled, an ordinary in-place update to such a leaf raises an error:

import torch

w = torch.tensor([1., 2.], requires_grad=True)
try:
    w.add_(0.1)
except RuntimeError:
    print("In-place leaf update rejected")

This restriction helps prevent ambiguous or invalid gradient histories. Being a leaf is about the recorded graph, not simply whether you typed a tensor constructor. Tensors that do not require gradients are also considered leaves by convention.

Update without recording the optimizer step

A conventional SGD update changes the parameter after its gradient has been computed. Use a no-grad block to make that update without recording it as part of the next differentiation graph:

w = torch.tensor([1., 2.], requires_grad=True)
grad = torch.tensor([0.5, -1.0])
with torch.no_grad():
    w.add_(-0.1 * grad)
print([round(v, 2) for v in w.tolist()])  # [0.95, 2.1]
print(w.is_leaf, w.requires_grad)  # True True

The block does not turn off the parameter's requires_grad flag. Future operations can still record gradients through it.

No-grad does not make every mutation safe

Existing graphs may have saved tensor values needed for backward. A mutation inside no-grad still changes those values and updates the tensor's version counter. Backward through an older graph can therefore raise a version-mismatch error. Finish the relevant backward pass before a conventional parameter update.

Ordinary optimizer steps typically avoid recording the update. That is not universal: differentiable optimization deliberately records update computations and requires a different design.

Preserve the caller's parameter

Writing w = w - lr * grad only rebinds the local Python name. It does not update the original tensor held by the caller or optimizer. This task requires changing that existing tensor while preserving its storage.

Gradient clearing is separate from the update. Setting w.grad = None and zeroing an existing gradient buffer can have different effects on optimizer behavior; for example, an optimizer may skip a parameter with no gradient. This function receives grad explicitly and should not clear it.

Your task

def sgd_step(w: torch.Tensor, grad: torch.Tensor, lr: float) -> None

The harness supplies a floating-point leaf w requiring gradients and a same-shaped gradient tensor. Update w in place to w - lr * grad and return None. Preserve its storage, leaf status and requires_grad=True. The supplied gradient has already been computed.