Lightning-AI/pytorch-lightning · error · KeyError

Trying to restore learning rate scheduler state but checkpoi

Error message

Trying to restore learning rate scheduler state but checkpoint contains only the model. This is probably due to `ModelCheckpoint.save_weights_only` being set to `True`.

What it means

Companion to the optimizer error: raised when the strategy restores training state and the checkpoint is missing the 'lr_schedulers' key, i.e., it was saved with ModelCheckpoint(save_weights_only=True). Learning-rate scheduler state cannot be restored from a weights-only file.

Source

Thrown at src/lightning/pytorch/trainer/connectors/checkpoint_connector.py:382

                f" but you have set Trainer(max_epochs={self.trainer.max_epochs})."
            )

    def restore_optimizers_and_schedulers(self) -> None:
        """Restores the optimizers and learning rate scheduler states from the pre-loaded checkpoint."""
        if not self._loaded_checkpoint:
            return

        if self.trainer.strategy.lightning_restore_optimizer:
            # validation
            if "optimizer_states" not in self._loaded_checkpoint:
                raise KeyError(
                    "Trying to restore optimizer state but checkpoint contains only the model."
                    " This is probably due to `ModelCheckpoint.save_weights_only` being set to `True`."
                )
            self.restore_optimizers()

        if "lr_schedulers" not in self._loaded_checkpoint:
            raise KeyError(
                "Trying to restore learning rate scheduler state but checkpoint contains only the model."
                " This is probably due to `ModelCheckpoint.save_weights_only` being set to `True`."
            )
        self.restore_lr_schedulers()

    def restore_optimizers(self) -> None:
        """Restores the optimizer states from the pre-loaded checkpoint."""
        if not self._loaded_checkpoint:
            return

        # restore the optimizers
        self.trainer.strategy.load_optimizer_state_dict(self._loaded_checkpoint)

    def restore_lr_schedulers(self) -> None:
        """Restores the learning rate scheduler states from the pre-loaded checkpoint."""
        if not self._loaded_checkpoint:
            return

View on GitHub (pinned to 9fed5c27d2)

Solutions

1. Save checkpoints with save_weights_only=False to include lr_schedulers and optimizer states
2. Load weights only and initialize a new scheduler (fresh training loop) instead of resuming
3. Resume from a complete checkpoint file that was saved with full training state

Example fix

# before
trainer.fit(model, ckpt_path="weights_only.ckpt")  # KeyError: lr_schedulers
# after
ModelCheckpoint(save_weights_only=False)  # during original fit
trainer.fit(model, ckpt_path="full_state.ckpt")

import torch
ckpt = torch.load(path, map_location="cpu", weights_only=False)
assert "lr_schedulers" in ckpt and "optimizer_states" in ckpt, "weights-only checkpoint cannot resume scheduler state"

When it happens

Trigger: Same as 552: trainer.fit(model, ckpt_path=<weights-only checkpoint>) where lightning_restore_optimizer applies; the optimizer key may exist if schedulers alone were excluded, but typically both are missing.

Common situations: Resuming from artifacts saved for inference; HuggingFace-style released checkpoints; switching checkpoint config between experiments.

Related errors

• Trying to restore optimizer state but checkpoint contains on
• `.{fn}(ckpt_path="hpc")` is set but no HPC checkpoint was fo
• `.{fn}()` found no path for the best weights: {ckpt_path!r}.
• Could not find a distributed model in the provided checkpoin
• Found multiple distributed models in the given state. Loadin

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