def test_lr_scheduler_step_hook(tmp_path):
    """Test that custom lr scheduler works and `lr_scheduler_step` is called at appropriate time."""

    class CustomEpochScheduler:
        def __init__(self, optimizer):
            self.optimizer = optimizer

        def step(self, epoch): ...

        def state_dict(self): ...

        def load_state_dict(self, state_dict): ...

    class CustomBoringModel(BoringModel):
        def lr_scheduler_step(self, scheduler: int, metric):
            # step-level
            if isinstance(scheduler, torch.optim.lr_scheduler.StepLR):
                super().lr_scheduler_step(scheduler, metric)
            # epoch-level, custom scheduler
            elif isinstance(scheduler, CustomEpochScheduler):
                scheduler.step(epoch=self.current_epoch)

        def configure_optimizers(self):
            optimizer = torch.optim.SGD(self.layer.parameters(), lr=1e-2)
            lr_scheduler1 = {"scheduler": torch.optim.lr_scheduler.StepLR(optimizer, step_size=1), "interval": "step"}
            lr_scheduler2 = CustomEpochScheduler(optimizer)
            return [optimizer], [lr_scheduler1, lr_scheduler2]

    model = CustomBoringModel()
    max_epochs = 3
    limit_train_batches = 2
    trainer = Trainer(
        default_root_dir=tmp_path,
        enable_checkpointing=False,
        logger=False,
        max_epochs=max_epochs,
        limit_train_batches=limit_train_batches,
        limit_val_batches=0,
    )
    with (
        mock.patch.object(CustomEpochScheduler, "step") as mock_method_epoch,
        mock.patch.object(torch.optim.lr_scheduler.StepLR, "step") as mock_method_step,
    ):
        trainer.fit(model)

    assert mock_method_epoch.mock_calls == [call(epoch=e) for e in range(max_epochs)]
    # first step is called by PyTorch LRScheduler
    assert mock_method_step.call_count == max_epochs * limit_train_batches + 1
Test that custom lr scheduler works and `lr_scheduler_step` is called at appropriate time.
tests/tests_pytorch/trainer/optimization/test_optimizers.py