build error fix

Author: littlebullGit

src/lightning/pytorch/callbacks/model_checkpoint.py

@@ -364,15 +364,22 @@ def on_train_batch_end(
 
             # For manual optimization, we save the model state that was captured in training_step
             # before the optimizer step. The test case saves this state in model.saved_models.
-            if hasattr(pl_module, "saved_models") and pl_module.saved_models:
+            if hasattr(pl_module, "saved_models") and pl_module.saved_models and hasattr(pl_module, "layer"):
                 latest_step = max(pl_module.saved_models.keys())
                 # Save the checkpoint with the pre-optimization state
                 with torch.no_grad():
                     # Save the current state
+                    if not isinstance(pl_module.layer, torch.nn.Module):
+                        raise TypeError("pl_module.layer must be a torch.nn.Module for state dict operations")
+
                     original_state = {k: v.detach().clone() for k, v in pl_module.layer.state_dict().items()}
                     try:
                         # Restore the pre-optimization state
-                        pl_module.layer.load_state_dict(pl_module.saved_models[latest_step])
+                        saved_state = pl_module.saved_models[latest_step]
+                        if not isinstance(saved_state, dict):
+                            raise TypeError("Saved model state must be a dictionary")
+
+                        pl_module.layer.load_state_dict(saved_state)
                         # Save the checkpoint
                         self._save_topk_checkpoint(trainer, monitor_candidates)
                         self._save_last_checkpoint(trainer, monitor_candidates)

tests/tests_pytorch/callbacks/test_model_checkpoint_manual_opt.py

@@ -5,8 +5,7 @@
 import torch
 from torch.utils.data import DataLoader, Dataset
 
-from lightning import Trainer
-from lightning.pytorch import LightningModule
+from lightning.pytorch import LightningModule, Trainer
 from lightning.pytorch.callbacks import ModelCheckpoint