Repository: Lightning-AI/pytorch-lightning
File: tests/tests_pytorch/tuner/test_lr_finder.py (Python)

Code before:
            return loss

        def configure_optimizers(self):
            return torch.optim.Adam(self.parameters(), lr=self.hparams.lr)

    seed_everything(123)

    ds = RandomDataset(n=512)
    train_loader = DataLoader(ds, batch_size=64, shuffle=False)

    model = TinyAE(lr=1e-5)

    lr_finder_cb = LearningRateFinder()  # default update_attr=True should apply suggestion
    lr_monitor = LearningRateMonitor(logging_interval="step")

    trainer = Trainer(
        default_root_dir=tmp_path,
        max_epochs=2,
        callbacks=[lr_finder_cb, lr_monitor],
        enable_model_summary=False,
        enable_progress_bar=False,
        log_every_n_steps=1,
    )

    trainer.fit(model, train_loader)

    # Ensure LR Finder produced a suggestion for this setup; if not, the test can't assert application
    assert lr_finder_cb.optimal_lr is not None, "LR Finder should have computed results"
    suggestion = lr_finder_cb.optimal_lr.suggestion()
    assert suggestion is not None, "LR Finder should produce a suggestion for this setup"

    # Verify that the optimizer used for subsequent training has the suggested LR applied
    assert trainer.optimizers, "Trainer should have an optimizer after fit"
    current_lr = trainer.optimizers[0].param_groups[0]["lr"]
    assert current_lr == pytest.approx(suggestion), (
        f"LR Finder suggestion {suggestion} should be applied to optimizer, but got {current_lr}"
    )


def test_exponential_vs_linear_mode_gradient_difference(tmp_path):
    """Test that exponential and linear modes produce different but valid suggestions.

    This verifies that the spacing fix works for both modes and that they behave differently as expected due to their
    different lr progressions.

    """

    class TestModel(BoringModel):
        def __init__(self):
            super().__init__()
            self.lr = 1e-3

Review comment (suggestion):
```suggestion
    assert model.hparams.lr is not None
```

Code after:

    seed_everything(123)

    ds = RandomDataset(n=512)
    train_loader = DataLoader(ds, batch_size=64, shuffle=False)

    model = TinyAE(lr=1e-5)

    lr_finder_cb = LearningRateFinder()  # default update_attr=True should apply suggestion
    lr_monitor = LearningRateMonitor(logging_interval="step")

    trainer = Trainer(
        default_root_dir=tmp_path,
        max_epochs=2,
        callbacks=[lr_finder_cb, lr_monitor],
        enable_model_summary=False,
        enable_progress_bar=False,
        log_every_n_steps=1,
    )

    trainer.fit(model, train_loader)
    assert model.hparams.lr is not None
    # Ensure LR Finder produced a suggestion for this setup; if not, the test can't assert application
    assert lr_finder_cb.optimal_lr is not None, "LR Finder should have computed results"
    suggestion = lr_finder_cb.optimal_lr.suggestion()
    assert suggestion is not None, "LR Finder should produce a suggestion for this setup"

    # Verify that the optimizer used for subsequent training has the suggested LR applied
    assert trainer.optimizers, "Trainer should have an optimizer after fit"
    current_lr = trainer.optimizers[0].param_groups[0]["lr"]
    assert current_lr == pytest.approx(suggestion), (
        f"LR Finder suggestion {suggestion} should be applied to optimizer, but got {current_lr}"
    )


def test_exponential_vs_linear_mode_gradient_difference(tmp_path):
    """Test that exponential and linear modes produce different but valid suggestions.

    This verifies that the spacing fix works for both modes and that they behave differently as expected due to their
    different lr progressions.

    """

    class TestModel(BoringModel):
        def __init__(self):
            super().__init__()
            self.lr = 1e-3

    seed_everything(42)

    # Test both modes with identical parameters