def test_scheduler_step_count(mock_step, tmp_path, max_epoch, limit_train_batches, interval):
    """Test to ensure that the scheduler is called the correct amount of times during training when scheduler is set to
    step or epoch."""

    class TestModel(BoringModel):
        def configure_optimizers(self):
            optimizer = torch.optim.SGD(self.layer.parameters(), lr=0.1)
            scheduler = torch.optim.lr_scheduler.StepLR(optimizer, 1, gamma=0.1)
            return {
                "optimizer": optimizer,
                "lr_scheduler": {"scheduler": scheduler, "interval": interval},
            }

    model = TestModel()
    trainer = Trainer(
        default_root_dir=tmp_path,
        limit_train_batches=limit_train_batches,
        limit_val_batches=0,
        max_epochs=max_epoch,
        accelerator="gpu",
        devices=1,
        strategy="deepspeed",
        enable_progress_bar=False,
        enable_model_summary=False,
    )
    trainer.fit(model)
    if interval == "epoch":
        # assert called once at init and once during training
        assert mock_step.call_count == 1 + max_epoch
    else:
        # assert called once at init and once during training
        assert mock_step.call_count == 1 + (max_epoch * limit_train_batches)
Test to ensure that the scheduler is called the correct amount of times during training when scheduler is set to
    step or epoch.
tests/tests_pytorch/strategies/test_deepspeed.py