Repository: Lightning-AI/pytorch-lightning
File: pytorch_lightning/callbacks/lr_monitor.py (Python)

Code before:
    def on_train_start(self, trainer, pl_module):
        """
        Called before training, determines unique names for all lr
        schedulers in the case of multiple of the same type or in
        the case of multiple parameter groups
        """
        if not trainer.logger:
            raise MisconfigurationException(
                'Cannot use LearningRateMonitor callback with Trainer that has no logger.'
            )

        if not trainer.lr_schedulers:
            rank_zero_warn(
                'You are using LearningRateMonitor callback with models that'
                ' have no learning rate schedulers. Please see documentation'
                ' for `configure_optimizers` method.', RuntimeWarning
            )

        # Find names for schedulers
        names = self._find_names(trainer.lr_schedulers)

        # Initialize for storing values
        self.lrs = {name: [] for name in names}

    def on_train_batch_start(self, trainer, pl_module, *_):
        if not LearningRateMonitor._should_log(trainer):
            return

        if self.logging_interval != 'epoch':
            interval = 'step' if self.logging_interval is None else 'any'
            latest_stat = self._extract_lr(trainer, interval)

            if trainer.logger is not None and latest_stat:
                trainer.logger.log_metrics(latest_stat, step=trainer.global_step)

    def on_train_epoch_start(self, trainer, pl_module):
        if self.logging_interval != 'step':
            interval = 'epoch' if self.logging_interval is None else 'any'
            latest_stat = self._extract_lr(trainer, interval)

            if trainer.logger is not None and latest_stat:
                trainer.logger.log_metrics(latest_stat, step=trainer.current_epoch)

    def _extract_lr(self, trainer, interval):
        latest_stat = {}

        for name, scheduler in zip(self.lr_sch_names, trainer.lr_schedulers):
            if scheduler['interval'] == interval or interval == 'any':
                param_groups = scheduler['scheduler'].optimizer.param_groups
                if len(param_groups) != 1:
                    for i, pg in enumerate(param_groups):

Review comment (suggestion):
```suggestion
        if not self._should_log(trainer):
```

Code after:
    def on_train_start(self, trainer, *args, **kwargs):
        """
        Called before training, determines unique names for all lr
        schedulers in the case of multiple of the same type or in
        the case of multiple parameter groups
        """
        if not trainer.logger:
            raise MisconfigurationException(
                'Cannot use LearningRateMonitor callback with Trainer that has no logger.'
            )

        if not trainer.lr_schedulers:
            rank_zero_warn(
                'You are using LearningRateMonitor callback with models that'
                ' have no learning rate schedulers. Please see documentation'
                ' for `configure_optimizers` method.', RuntimeWarning
            )

        # Find names for schedulers
        names = self._find_names(trainer.lr_schedulers)

        # Initialize for storing values
        self.lrs = {name: [] for name in names}

    def on_train_batch_start(self, trainer, *args, **kwargs):
        if not self._should_log(trainer):
            return

        if self.logging_interval != 'epoch':
            interval = 'step' if self.logging_interval is None else 'any'
            latest_stat = self._extract_lr(trainer, interval)

            if trainer.logger is not None and latest_stat:
                trainer.logger.log_metrics(latest_stat, step=trainer.global_step)

    def on_train_epoch_start(self, trainer, *args, **kwargs):
        if self.logging_interval != 'step':
            interval = 'epoch' if self.logging_interval is None else 'any'
            latest_stat = self._extract_lr(trainer, interval)

            if trainer.logger is not None and latest_stat:
                trainer.logger.log_metrics(latest_stat, step=trainer.current_epoch)

    def _extract_lr(self, trainer, interval: str) -> Dict[str, float]:
        latest_stat = {}

        for name, scheduler in zip(self.lr_sch_names, trainer.lr_schedulers):
            if scheduler['interval'] == interval or interval == 'any':
                param_groups = scheduler['scheduler'].optimizer.param_groups
                if len(param_groups) != 1:
                    for i, pg in enumerate(param_groups):