One learning rate is rarely right for a whole run. Early on, big steps cover ground fast; near the end, small steps settle into the minimum instead of bouncing around it. A scheduler changes the optimizer's rate as training goes on.

import torch

w = torch.zeros(1, requires_grad=True)
opt = torch.optim.SGD([w], lr=1.0)
sched = torch.optim.lr_scheduler.StepLR(
    opt, step_size=2, gamma=0.5)
for epoch in range(6):
    opt.step()  # after loss.backward()
    sched.step()
    print(epoch, sched.get_last_lr())

Call sched.step() once per epoch, or once per batch if the schedule counts batches, always after opt.step(). The current rate is opt.param_groups[0]["lr"].

A common recipe for bigger models is warm-up then cosine decay: the rate rises from 0 over the first steps, while early gradients are still chaotic, then falls along half a cosine to 0. LambdaLR builds any schedule from a function of the step number, whose value multiplies the rate the optimizer started with.

Write warmup_cosine(step, total, warmup, peak): the rate rises linearly from 0 at step 0 to peak at step warmup, then falls along half a cosine, peak * 0.5 * (1 + cos(pi * progress)), to 0 at step total, and stays 0 after it. Write make_scheduler(optimizer, total, warmup) returning a LambdaLR with this shape, the optimizer's own rate as the peak.

The explanation opens here once every check passes.