def test_lightning_optimizer(tmp_path):
    """Test that optimizer are correctly wrapped by our LightningOptimizer."""

    class TestModel(BoringModel):
        def configure_optimizers(self):
            optimizer = torch.optim.SGD(self.layer.parameters(), lr=0.1)
            lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=1)
            return [optimizer], [lr_scheduler]

    model = TestModel()
    trainer = Trainer(
        default_root_dir=tmp_path, limit_train_batches=1, limit_val_batches=1, max_epochs=1, enable_model_summary=False
    )
    trainer.fit(model)

    lightning_opt = model.optimizers()
    assert str(lightning_opt) == "Lightning" + str(lightning_opt.optimizer)
Test that optimizer are correctly wrapped by our LightningOptimizer.
tests/tests_pytorch/core/test_lightning_optimizer.py