When do you call a PyTorch lr_scheduler's step() — every batch or every epoch?

middleshould knowfreq 48%

answer

the counter has no units of its own
the schedule you chose decides the granularity
after the optimizer, not before
one scheduler wants a metric
print the rate; do not assume it moved

basics

~20 s

It depends on the schedule. Epoch-granularity schedules such as StepLR are stepped once per epoch; warmup and OneCycleLR are stepped once per optimizer step. Either way the call goes after optimizer.step(), and ReduceLROnPlateau also needs a metric argument.

solid answer

~50 s

A scheduler in torch.optim.lr_scheduler computes the learning rate from an internal counter, and step() is what advances that counter — so the granularity is a property of the schedule you chose, not a global rule. StepLR, MultiStepLR and a CosineAnnealingLR with T_max expressed in epochs are stepped once per epoch. OneCycleLR, LinearLR used as warmup, and any schedule whose horizon you defined in optimizer steps are stepped once per batch. Getting the granularity wrong does not raise — it just makes the decay finish in a fraction of the run or never arrive at all.

Two details are asked about constantly. First, scheduler.step() must come after optimizer.step(); PyTorch detects the reverse order and warns that the first learning-rate value of the schedule will be skipped. Second, ReduceLROnPlateau is not a formula-driven schedule — it takes the monitored metric as an argument, scheduler.step(val_loss), and is stepped once per validation, not per batch. Verify with scheduler.get_last_lr().

A scheduler in `torch.optim.lr_scheduler` computes the learning rate from an internal counter, and `step()` is what advances that counter — so the granularity is a property of the schedule you chose, not a global rule. `StepLR`, `MultiStepLR` and a `CosineAnnealingLR` with `T_max` expressed in epochs are stepped once per epoch. `OneCycleLR`, `LinearLR` used as warmup, and any schedule whose horizon you defined in optimizer steps are stepped once per batch. Getting the granularity wrong does not raise — it just makes the decay finish in a fraction of the run or never arrive at all. Two details are asked about constantly. First, `scheduler.step()` must come **after** `optimizer.step()`; PyTorch detects the reverse order and warns that the first learning-rate value of the schedule will be skipped. Second, `ReduceLROnPlateau` is not a formula-driven schedule — it takes the monitored metric as an argument, `scheduler.step(val_loss)`, and is stepped once per validation, not per batch. Verify with `scheduler.get_last_lr()`.

code

python · 14 lines

python

import torch
from torch import nn

model = nn.Linear(4, 2)
opt = torch.optim.SGD(model.parameters(), lr=0.1)
sched = torch.optim.lr_scheduler.CosineAnnealingLR(opt, T_max=5)

for epoch in range(5):
    for _ in range(3):                       # batches
        opt.zero_grad(set_to_none=True)
        model(torch.randn(8, 4)).pow(2).mean().backward()
        opt.step()
    sched.step()                             # epoch-granularity schedule
    print(epoch, sched.get_last_lr())

go deeper

for a junior

Know that a scheduler changes the learning rate over time, that you construct it around an optimizer, and that its step() call is separate from the optimizer's.

for a middle

Explain that the scheduler's counter is unitless and that the granularity comes from where you call it, name which built-ins are per-epoch and which per-step, and state the after-the-optimizer ordering rule.

for a senior

Compose warmup with decay via SequentialLR, keep the scheduler counter aligned with real optimizer steps under accumulation and AMP skips, and always verify the rate by logging it rather than trusting the code.

for a principal

Own the schedule as a design decision — total-step budget, warmup length and floor are chosen together with batch size — and make the learning-rate trace a standard panel on every training dashboard so silent misconfiguration cannot survive a run.

The contract

A PyTorch learning-rate scheduler wraps an optimizer and rewrites the lr entry of each of its param_groups. It holds a counter — last_epoch, whatever the units really are — and step() advances it and recomputes the rate. The scheduler has no idea whether one increment means a batch or an epoch; that meaning comes entirely from where you call it and from how you parameterized the schedule.

That single fact resolves most confusion. CosineAnnealingLR(optimizer, T_max=50) decays over 50 step() calls. Put that call at the end of each epoch and you get a 50-epoch cosine. Put it at the end of each batch in a run with 500 batches per epoch and the cosine completes one tenth of the way through the first epoch and then starts rising again. Nothing errors; the run just trains at a rate you did not intend.

Per-epoch versus per-step schedules

Per-epoch by convention: StepLR (multiply by gamma every step_size counter ticks), MultiStepLR (drop at listed milestones), ExponentialLR, and CosineAnnealingLR when T_max is a number of epochs.

Per-optimizer-step: OneCycleLR, which is explicitly designed around a total step count and is documented to be stepped after every batch; LinearLR or LambdaLR used to implement warmup, since warmup is measured in hundreds of steps, not epochs; and any cosine whose T_max you set to the total number of training steps — which is the standard practice in transformer training, precisely because per-step decay is smoother.

OneCycleLR is the one that fails loudly: it knows its total budget, either from total_steps or from steps_per_epoch times epochs, and it raises a ValueError if you step it more times than that. Treat that error as a signal that your step accounting is off — most often because gradient accumulation means the number of optimizer steps is a fraction of the number of batches.

ReduceLROnPlateau is different

ReduceLROnPlateau is not a formula of the counter; it watches a metric. You call scheduler.step(metric) with the value you are monitoring — typically the validation loss — once per evaluation. If the metric has failed to improve for patience calls, the rate is multiplied by factor. mode='min' or 'max' tells it which direction counts as improvement, and threshold sets how much of a change counts at all.

Because it consumes a metric, it is the one scheduler where calling step() with no argument is a mistake, and where per-batch stepping is nearly always wrong: patience would be counted in batches, and the noise in a single batch's loss makes the plateau detection meaningless.

Order relative to optimizer.step()

The learning rate that applies to an update is the one in param_groups at the moment optimizer.step() runs. So the ordering within an iteration is: optimizer.step() first, then scheduler.step() to prepare the rate for the next update. PyTorch actively checks for the inverted order and emits a warning saying that lr_scheduler.step() was detected before optimizer.step() and that the first value of the schedule will be skipped. It is a warning, not an error, and it is routinely ignored to no great harm — but it is a giveaway that the author copied the loop rather than reasoned about it.

The same ordering point interacts with mixed precision: scaler.step(optimizer) may skip the update entirely when it finds non-finite gradients, while scheduler.step() advances regardless. Over a long run with a few skipped steps this drift is negligible, but it means "scheduler counter" and "number of applied updates" are not exactly equal.

Composing warmup with decay

The common recipe — linear warmup for a few hundred steps, then cosine decay — is built with SequentialLR, which takes a list of schedulers and a list of milestone counts and hands control from one to the next: LinearLR for the warmup phase, then CosineAnnealingLR. ChainedScheduler is the different thing that composes schedulers multiplicatively at every step rather than switching between them, and LambdaLR lets you write the whole curve as a function of the step index when the built-ins do not fit.

Verifying it works

Do not assume. Print scheduler.get_last_lr() (a list, one entry per parameter group) or read optimizer.param_groups[0]["lr"] every few hundred steps and log it alongside the loss. The two failures this catches immediately are a schedule that never moves — because step() is called on the optimizer only — and a schedule that finished in the first epoch because of a granularity mismatch. And remember the scheduler's state_dict() belongs in your checkpoints, or a resumed run restarts the curve from the beginning.

## The contract A PyTorch learning-rate scheduler wraps an optimizer and rewrites the `lr` entry of each of its `param_groups`. It holds a counter — `last_epoch`, whatever the units really are — and `step()` advances it and recomputes the rate. The scheduler has no idea whether one increment means a batch or an epoch; that meaning comes entirely from where you call it and from how you parameterized the schedule. That single fact resolves most confusion. `CosineAnnealingLR(optimizer, T_max=50)` decays over 50 `step()` calls. Put that call at the end of each epoch and you get a 50-epoch cosine. Put it at the end of each batch in a run with 500 batches per epoch and the cosine completes one tenth of the way through the first epoch and then starts rising again. Nothing errors; the run just trains at a rate you did not intend. ## Per-epoch versus per-step schedules **Per-epoch by convention**: `StepLR` (multiply by `gamma` every `step_size` counter ticks), `MultiStepLR` (drop at listed milestones), `ExponentialLR`, and `CosineAnnealingLR` when `T_max` is a number of epochs. **Per-optimizer-step**: `OneCycleLR`, which is explicitly designed around a total step count and is documented to be stepped after every batch; `LinearLR` or `LambdaLR` used to implement warmup, since warmup is measured in hundreds of steps, not epochs; and any cosine whose `T_max` you set to the total number of training steps — which is the standard practice in transformer training, precisely because per-step decay is smoother. `OneCycleLR` is the one that fails loudly: it knows its total budget, either from `total_steps` or from `steps_per_epoch` times `epochs`, and it raises a `ValueError` if you step it more times than that. Treat that error as a signal that your step accounting is off — most often because gradient accumulation means the number of optimizer steps is a fraction of the number of batches. ## ReduceLROnPlateau is different `ReduceLROnPlateau` is not a formula of the counter; it watches a metric. You call `scheduler.step(metric)` with the value you are monitoring — typically the validation loss — once per evaluation. If the metric has failed to improve for `patience` calls, the rate is multiplied by `factor`. `mode='min'` or `'max'` tells it which direction counts as improvement, and `threshold` sets how much of a change counts at all. Because it consumes a metric, it is the one scheduler where calling `step()` with no argument is a mistake, and where per-batch stepping is nearly always wrong: patience would be counted in batches, and the noise in a single batch's loss makes the plateau detection meaningless. ## Order relative to optimizer.step() The learning rate that applies to an update is the one in `param_groups` at the moment `optimizer.step()` runs. So the ordering within an iteration is: `optimizer.step()` first, then `scheduler.step()` to prepare the rate for the *next* update. PyTorch actively checks for the inverted order and emits a warning saying that `lr_scheduler.step()` was detected before `optimizer.step()` and that the first value of the schedule will be skipped. It is a warning, not an error, and it is routinely ignored to no great harm — but it is a giveaway that the author copied the loop rather than reasoned about it. The same ordering point interacts with mixed precision: `scaler.step(optimizer)` may skip the update entirely when it finds non-finite gradients, while `scheduler.step()` advances regardless. Over a long run with a few skipped steps this drift is negligible, but it means "scheduler counter" and "number of applied updates" are not exactly equal. ## Composing warmup with decay The common recipe — linear warmup for a few hundred steps, then cosine decay — is built with `SequentialLR`, which takes a list of schedulers and a list of milestone counts and hands control from one to the next: `LinearLR` for the warmup phase, then `CosineAnnealingLR`. `ChainedScheduler` is the different thing that composes schedulers multiplicatively at every step rather than switching between them, and `LambdaLR` lets you write the whole curve as a function of the step index when the built-ins do not fit. ## Verifying it works Do not assume. Print `scheduler.get_last_lr()` (a list, one entry per parameter group) or read `optimizer.param_groups[0]["lr"]` every few hundred steps and log it alongside the loss. The two failures this catches immediately are a schedule that never moves — because `step()` is called on the optimizer only — and a schedule that finished in the first epoch because of a granularity mismatch. And remember the scheduler's `state_dict()` belongs in your checkpoints, or a resumed run restarts the curve from the beginning.

How do you compose a linear warmup followed by cosine decay in PyTorch?
Build a LinearLR for the warmup and a CosineAnnealingLR for the decay, then wrap them in SequentialLR with a milestone at the step where control should hand over. Step the SequentialLR once per optimizer step, and size the cosine's T_max to the remaining steps rather than the total, or the curve will not reach its floor.
You suspect the learning rate is not changing at all. How do you check?
Log scheduler.get_last_lr(), which returns one value per parameter group, or read optimizer.param_groups[0]["lr"] directly, every few hundred iterations. A flat trace means step() is never being reached; a curve that flattens in epoch one means the schedule's horizon was defined in epochs but stepped per batch.
With gradient accumulation, does the scheduler step per batch or per optimizer step?
Per optimizer step — the scheduler should advance only when an update was actually applied, so it belongs inside the same branch as optimizer.step(). Getting this wrong is the usual cause of OneCycleLR raising a ValueError about exceeding its total step budget, since the batch count is a multiple of the real update count.

saying these in an interview costs you the question

Calling scheduler.step() before optimizer.step()
Stepping an epoch-granularity schedule once per batch
Calling ReduceLROnPlateau.step() with no metric argument
Assuming every scheduler is stepped at the same granularity
Leaving the scheduler state out of the checkpoint

When do you call a PyTorch lr_scheduler's step() — every batch or every epoch?

§ answer

basics

solid answer

code

§ go deeper

▾level notes4 levels

▾deep dive~4 min

The contract

Per-epoch versus per-step schedules

ReduceLROnPlateau is different

Order relative to optimizer.step()

Composing warmup with decay

Verifying it works

▾common follow-ups3 questions

▾red flags5 flags

§ context

▾related questions3 questions

▾related topics3 topics

The contract

Per-epoch versus per-step schedules

ReduceLROnPlateau is different

Order relative to optimizer.step()

Composing warmup with decay

Verifying it works

answer

go deeper

context