Lightning-AI/pytorch-lightning · error · TypeError

The provided lr scheduler `{scheduler.__class__.__name__}` i

Error message

The provided lr scheduler `{scheduler.__class__.__name__}` is invalid. It should have `state_dict` and `load_state_dict` methods defined.

What it means

Every scheduler Lightning manages must be checkpointable, i.e. implement state_dict/load_state_dict (the _Stateful protocol). A scheduler object lacking these methods raises TypeError during optimizer/scheduler setup.

Source

Thrown at src/lightning/pytorch/core/optimizer.py:334

            if keys_to_warn:
                rank_zero_warn(
                    f"The lr scheduler dict contains the key(s) {keys_to_warn}, but the keys will be ignored."
                    " You need to call `lr_scheduler.step()` manually in manual optimization.",
                    category=RuntimeWarning,
                )

            config = LRSchedulerConfig(**{key: scheduler[key] for key in scheduler if key not in invalid_keys})
        else:
            config = LRSchedulerConfig(scheduler)
        lr_scheduler_configs.append(config)
    return lr_scheduler_configs


def _validate_scheduler_api(lr_scheduler_configs: list[LRSchedulerConfig], model: "pl.LightningModule") -> None:
    for config in lr_scheduler_configs:
        scheduler = config.scheduler
        if not isinstance(scheduler, _Stateful):
            raise TypeError(
                f"The provided lr scheduler `{scheduler.__class__.__name__}` is invalid."
                " It should have `state_dict` and `load_state_dict` methods defined."
            )

        if (
            not isinstance(scheduler, LRSchedulerTypeTuple)
            and not is_overridden("lr_scheduler_step", model)
            and model.automatic_optimization
        ):
            raise MisconfigurationException(
                f"The provided lr scheduler `{scheduler.__class__.__name__}` doesn't follow PyTorch's LRScheduler"
                " API. You should override the `LightningModule.lr_scheduler_step` hook with your own logic if"
                " you are using a custom LR scheduler."
            )


def _validate_multiple_optimizers_support(optimizers: list[Optimizer], model: "pl.LightningModule") -> None:
    if is_param_in_hook_signature(model.training_step, "optimizer_idx", explicit=True):

View on GitHub (pinned to 9fed5c27d2)

Solutions

  1. Subclass torch.optim.lr_scheduler.LRScheduler (which provides both methods)
  2. Or implement state_dict() and load_state_dict(state_dict) on the custom scheduler class

Example fix

# before
class MySched:
    def __init__(self, opt): self.opt = opt
    def step(self): ...
# after
class MySched(torch.optim.lr_scheduler.LRScheduler):
    def get_lr(self):
        return [g['lr'] for g in self.optimizer.param_groups]
Defensive patterns

Strategy: type-guard

Validate before calling

assert hasattr(scheduler, "state_dict") and hasattr(scheduler, "load_state_dict")

Type guard

def is_stateful(sched) -> bool:
    return hasattr(sched, "state_dict") and hasattr(sched, "load_state_dict")

Prevention

When it happens

Trigger: Returning a custom scheduler class that does not subclass torch.optim.lr_scheduler.LRScheduler and does not implement state_dict/load_state_dict.

Common situations: Hand-rolled warmup or custom LR wrappers that forget serialization methods needed for checkpoint resume.

Related errors


AI-assisted analysis of Lightning-AI/pytorch-lightning@9fed5c27d2 (2026-08-28). Data as JSON: /api/errors/e603e6e4906ace67. Report an issue: GitHub.