Repository: Lightning-AI/pytorch-lightning
File: pytorch_lightning/trainer/connectors/checkpoint_connector.py (Python)

Code before:
            return

        # restore the optimizers
        optimizer_states = checkpoint['optimizer_states']
        for optimizer, opt_state in zip(self.trainer.optimizers, optimizer_states):
            optimizer.load_state_dict(opt_state)

            # move optimizer to GPU 1 weight at a time
            # avoids OOM
            if self.trainer.root_gpu is not None:
                for state in optimizer.state.values():
                    for k, v in state.items():
                        if isinstance(v, torch.Tensor):
                            state[k] = v.cuda(self.trainer.root_gpu)

        # restore the lr schedulers
        lr_schedulers = checkpoint['lr_schedulers']
        for scheduler, lrs_state in zip(self.trainer.lr_schedulers, lr_schedulers):
            scheduler['scheduler'].load_state_dict(lrs_state)

    def restore_callbacks(self) -> None:
        """ Restores all callbacks from the pre-loaded checkpoint. """
        if not self._loaded_checkpoint:
            return

        if any([key in self._loaded_checkpoint for key in DEPRECATED_CHECKPOINT_KEYS]):
            raise ValueError(
                "The checkpoint you're attempting to load follows an"
                " outdated schema. You can upgrade to the current schema by running"
                " `python -m pytorch_lightning.utilities.upgrade_checkpoint --file model.ckpt`"
                " where `model.ckpt` is your checkpoint file."
            )
        self.trainer.on_load_checkpoint(self._loaded_checkpoint)

    def restore_progress(self) -> None:
        """
        Restores the training progress from the pre-loaded checkpoint. This currently includes only the global step
        and current epoch.
        """
        if not self._loaded_checkpoint:
            return

        self.trainer.train_loop.global_step = self._loaded_checkpoint['global_step']
        self.trainer.train_loop.current_epoch = self._loaded_checkpoint['epoch']

        # crash if max_epochs is lower then the current epoch from the checkpoint
        if 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 the Trainer(max_epochs={self.trainer.max_epochs})"
            )

Review comment (suggestion):
```suggestion
        if any(key in self._loaded_checkpoint for key in DEPRECATED_CHECKPOINT_KEYS):
```

Code after:
            return

        # restore the optimizers
        optimizer_states = checkpoint['optimizer_states']
        for optimizer, opt_state in zip(self.trainer.optimizers, optimizer_states):
            optimizer.load_state_dict(opt_state)

            # move optimizer to GPU 1 weight at a time
            # avoids OOM
            if self.trainer.root_gpu is not None:
                for state in optimizer.state.values():
                    for k, v in state.items():
                        if isinstance(v, torch.Tensor):
                            state[k] = v.cuda(self.trainer.root_gpu)

        # restore the lr schedulers
        lr_schedulers = checkpoint['lr_schedulers']
        for scheduler, lrs_state in zip(self.trainer.lr_schedulers, lr_schedulers):
            scheduler['scheduler'].load_state_dict(lrs_state)

    def restore_callbacks(self) -> None:
        """ Restores all callbacks from the pre-loaded checkpoint. """
        if not self._loaded_checkpoint:
            return

        if any(key in self._loaded_checkpoint for key in DEPRECATED_CHECKPOINT_KEYS):
            raise ValueError(
                "The checkpoint you're attempting to load follows an"
                " outdated schema. You can upgrade to the current schema by running"
                " `python -m pytorch_lightning.utilities.upgrade_checkpoint --file model.ckpt`"
                " where `model.ckpt` is your checkpoint file."
            )
        self.trainer.on_load_checkpoint(self._loaded_checkpoint)

    def restore_progress(self) -> None:
        """
        Restores the training progress from the pre-loaded checkpoint. This currently includes only the global step
        and current epoch.
        """
        if not self._loaded_checkpoint:
            return

        self.trainer.train_loop.global_step = self._loaded_checkpoint['global_step']
        self.trainer.train_loop.current_epoch = self._loaded_checkpoint['epoch']

        # crash if max_epochs is lower then the current epoch from the checkpoint
        if 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})."
            )