fix mypy error

Author: littlebullGit

src/lightning/pytorch/callbacks/model_checkpoint.py

@@ -364,18 +364,26 @@ 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 and hasattr(pl_module, "layer"):
-                latest_step = max(pl_module.saved_models.keys())
+            if (
+                hasattr(pl_module, "saved_models")
+                and isinstance(pl_module.saved_models, dict)
+                and pl_module.saved_models
+                and hasattr(pl_module, "layer")
+                and isinstance(pl_module.layer, torch.nn.Module)
+            ):
+                # Get the latest saved state
+                saved_models = pl_module.saved_models
+                if not saved_models:  # Check if dictionary is not empty
+                    return
+
+                latest_step = max(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
-                        saved_state = pl_module.saved_models[latest_step]
+                        saved_state = saved_models[latest_step]
                         if not isinstance(saved_state, dict):
                             raise TypeError("Saved model state must be a dictionary")