Fix not running test codes (#13089)

Co-authored-by: Carlos Mocholí <[email protected]>

Author: 2 people

tests/strategies/test_deepspeed_strategy.py

@@ -749,9 +749,7 @@ def test_deepspeed_multigpu_stage_3_resume_training(tmpdir):
     initial_trainer.fit(initial_model, datamodule=dm)
 
     class TestCallback(Callback):
-        def on_train_batch_start(
-            self, trainer: Trainer, pl_module: LightningModule, batch: Any, batch_idx: int
-        ) -> None:
+        def on_train_epoch_start(self, trainer: Trainer, pl_module: LightningModule) -> None:
             original_deepspeed_strategy = initial_trainer.strategy
             current_deepspeed_strategy = trainer.strategy
 
@@ -778,10 +776,12 @@ def on_train_batch_start(
     model = ModelParallelClassificationModel()
     trainer = Trainer(
         default_root_dir=tmpdir,
-        fast_dev_run=True,
         strategy=DeepSpeedStrategy(stage=3),
         accelerator="gpu",
         devices=1,
+        max_epochs=2,
+        limit_train_batches=1,
+        limit_val_batches=0,
         precision=16,
         callbacks=TestCallback(),
         enable_progress_bar=False,