def test_lr_monitor_single_lr_with_momentum(tmp_path, opt: str):
    """Test that learning rates, momentum and weight decay are extracted and logged for single lr scheduler."""

    class LogMomentumModel(BoringModel):
        def __init__(self, opt):
            super().__init__()
            self.opt = opt

        def configure_optimizers(self):
            if self.opt == "SGD":
                opt_kwargs = {"momentum": 0.9}
            elif self.opt == "Adam":
                opt_kwargs = {"betas": (0.9, 0.999)}

            optimizer = getattr(optim, self.opt)(self.parameters(), lr=1e-2, **opt_kwargs)
            lr_scheduler = optim.lr_scheduler.OneCycleLR(optimizer, max_lr=1e-2, total_steps=10_000)
            return [optimizer], [lr_scheduler]

    model = LogMomentumModel(opt=opt)
    lr_monitor = LearningRateMonitor(log_momentum=True, log_weight_decay=True)
    trainer = Trainer(
        default_root_dir=tmp_path,
        max_epochs=2,
        limit_val_batches=2,
        limit_train_batches=5,
        log_every_n_steps=1,
        callbacks=[lr_monitor],
        logger=CSVLogger(tmp_path),
    )
    trainer.fit(model)

    assert all(v is not None for v in lr_monitor.last_momentum_values.values()), "Expected momentum to be logged"
    assert len(lr_monitor.last_momentum_values) == len(trainer.lr_scheduler_configs)
    assert all(k == f"lr-{opt}-momentum" for k in lr_monitor.last_momentum_values)

    assert all(v is not None for v in lr_monitor.last_weight_decay_values.values()), (
        "Expected weight decay to be logged"
    )
    assert len(lr_monitor.last_weight_decay_values) == len(trainer.lr_scheduler_configs)
    assert all(k == f"lr-{opt}-weight_decay" for k in lr_monitor.last_weight_decay_values)
Test that learning rates, momentum and weight decay are extracted and logged for single lr scheduler.
tests/tests_pytorch/callbacks/test_lr_monitor.py