def on_train_batch_end(
        self, trainer: "pl.Trainer", pl_module: "pl.LightningModule", outputs: STEP_OUTPUT, batch: Any, batch_idx: int
    ) -> None:
        """Called when the training batch ends, logs the calculated loss."""
        if (trainer.fit_loop.batch_idx + 1) % trainer.accumulate_grad_batches != 0:
            return

        # _AutomaticOptimization.run turns None STEP_OUTPUT into an empty dict
        if not outputs:
            # need to add an element, because we also added one element to lrs in on_train_batch_start
            # so add nan, because they are not considered when computing the suggestion
            self.losses.append(float("nan"))
            return

        if self.progress_bar:
            self.progress_bar.update()

        loss_tensor = outputs if isinstance(outputs, torch.Tensor) else outputs["loss"]
        assert loss_tensor is not None
        current_loss = loss_tensor.item()
        current_step = trainer.global_step

        # Avg loss (loss with momentum) + smoothing
        self.avg_loss = self.beta * self.avg_loss + (1 - self.beta) * current_loss
        smoothed_loss = self.avg_loss / (1 - self.beta ** (current_step + 1))

        # Check if we diverging
        if (
            self.early_stop_threshold is not None
            and current_step > 1
            and smoothed_loss > self.early_stop_threshold * self.best_loss
        ):
            trainer.should_stop = True  # stop signal
            if self.progress_bar:
                self.progress_bar.close()

        trainer.should_stop = trainer.strategy.broadcast(trainer.should_stop)

        # Save best loss for diverging checking
        if smoothed_loss < self.best_loss or current_step == 1:
            self.best_loss = smoothed_loss

        self.losses.append(smoothed_loss)
Called when the training batch ends, logs the calculated loss.
src/lightning/pytorch/tuner/lr_finder.py