update assertion and add comment

Lightning-AI · Jan 4, 2022 · 26766bb · 26766bb
1 parent a8ed671
commit 26766bb
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/tests/strategies/test_deepspeed_strategy.py b/tests/strategies/test_deepspeed_strategy.py
@@ -558,8 +558,9 @@ def on_load_checkpoint(self, checkpoint: Dict[str, Any]) -> None:
         if not hasattr(self, "model"):
             self.configure_sharded_model()
 
+        # Lightning saves the lr schedulers, but DeepSpeed saves the optimizer states separately
         assert len(checkpoint["lr_schedulers"]) == 1
-        assert len(checkpoint["optimizer_states"]) == 1
+        assert len(checkpoint["optimizer_states"]) == 0
 
 
 class ManualModelParallelClassificationModel(ModelParallelClassificationModel):