def test_lr_monitor_multiple_param_groups_no_lr_scheduler(tmp_path):
    """Test that the `LearningRateMonitor` is able to log correct keys with multiple param groups and no
    lr_scheduler."""

    class TestModel(BoringModel):
        def __init__(self, lr, momentum):
            super().__init__()
            self.save_hyperparameters()
            self.linear_a = torch.nn.Linear(32, 16)
            self.linear_b = torch.nn.Linear(16, 2)

        def forward(self, x):
            x = self.linear_a(x)
            x = self.linear_b(x)
            return x

        def configure_optimizers(self):
            param_groups = [
                {
                    "params": list(self.linear_a.parameters()),
                    "weight_decay": 0.1,
                },
                {
                    "params": list(self.linear_b.parameters()),
                    "weight_decay": 0.1,
                },
            ]
            return torch.optim.Adam(param_groups, lr=self.hparams.lr, betas=self.hparams.momentum)

    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=2,
        callbacks=[lr_monitor],
        logger=CSVLogger(tmp_path),
        enable_progress_bar=False,
        enable_model_summary=False,
    )

    lr = 1e-2
    momentum = 0.7
    weight_decay = 0.1
    model = TestModel(lr=lr, momentum=(momentum, 0.999))
    trainer.fit(model)

    assert len(lr_monitor.lrs) == len(trainer.optimizers[0].param_groups)
    assert list(lr_monitor.lrs) == ["lr-Adam/pg1", "lr-Adam/pg2"]
    assert list(lr_monitor.last_momentum_values) == ["lr-Adam/pg1-momentum", "lr-Adam/pg2-momentum"]
    assert all(val == momentum for val in lr_monitor.last_momentum_values.values())
    assert list(lr_monitor.last_weight_decay_values) == ["lr-Adam/pg1-weight_decay", "lr-Adam/pg2-weight_decay"]
    assert all(val == weight_decay for val in lr_monitor.last_weight_decay_values.values())
    assert all(all(val == lr for val in lr_monitor.lrs[lr_key]) for lr_key in lr_monitor.lrs)
Test that the `LearningRateMonitor` is able to log correct keys with multiple param groups and no
    lr_scheduler.
tests/tests_pytorch/callbacks/test_lr_monitor.py