Lightning-AI/pytorch-lightning · error · KeyError

Trying to restore optimizer state but checkpoint contains on

Error message

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

What it means

Raised as KeyError in restore_optimizers_and_schedulers when resuming training state: the strategy wants to restore optimizer state (lightning_restore_optimizer true) but the loaded checkpoint lacks the 'optimizer_states' key. This happens when the checkpoint was saved with weights only.

Source

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

        if (
            self.trainer.max_epochs != -1
            and self.trainer.max_epochs is not None
            and self.trainer.current_epoch > self.trainer.max_epochs
        ):
            raise MisconfigurationException(
                f"You restored a checkpoint with current_epoch={self.trainer.current_epoch},"
                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

View on GitHub (pinned to 9fed5c27d2)

Solutions

1. Re-save or originally save with ModelCheckpoint(save_weights_only=False) so optimizer states are included
2. Start fresh optimizers: load only weights (model.load_from_checkpoint) and call fit without ckpt_path resume, or use a strategy/setup with lightning_restore_optimizer=False
3. If the full checkpoint exists elsewhere, resume from it instead

Example fix

# before
ckpt = ModelCheckpoint(save_weights_only=True)  # later: fit(ckpt_path=...) -> KeyError
# after
ckpt = ModelCheckpoint(save_weights_only=False)
trainer.fit(model, ckpt_path=ckpt.best_model_path)  # now restores optimizer state

import torch
ckpt = torch.load(path, map_location="cpu", weights_only=False)
has_full_state = "optimizer_states" in ckpt
if trainer.strategy.lightning_restore_optimizer and not has_full_state:
    # avoid KeyError: load weights only and start fresh
    model = type(model).load_from_checkpoint(path)
    trainer.fit(model)  # no ckpt_path

When it happens

Trigger: ModelCheckpoint(save_weights_only=True) during the original fit (or load_from_checkpoint-produced weights), then trainer.fit(model, ckpt_path=that_file) with a strategy that restores optimizers (default single-device/DDL; false for e.g. some fault-tolerant flows).

Common situations: Saving compact checkpoints to save disk and later deciding to resume training from them; downloading a released weights-only .ckpt; upgrading from versions where save_weights_only defaults or semantics differed.

Related errors

• Trying to restore learning rate scheduler state but checkpoi
• `.{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/d9c7c419f28503c5. Report an issue: GitHub.