class_docstr: Special callback used by the learning rate finder. This callback logs the learning rate before each batch and
  logs the corresponding loss after each batch.
  
  Args:
      num_training: number of iterations done by the learning rate finder
      early_stop_threshold: threshold for stopping the search. If the
          loss at any point is larger than ``early_stop_threshold*best_loss``
          then the search is stopped. To disable, set to ``None``.
      progress_bar_refresh_rate: rate to refresh the progress bar for
          the learning rate finder
      beta: smoothing value, the loss being logged is a running average of
          loss values logged until now. ``beta`` controls the forget rate i.e.
          if ``beta=0`` all past information is ignored.
docstr: Called when the training batch ends, logs the calculated loss.
docstr_tok: ["Called","when","the","training","batch","ends",",","logs","the","calculated","loss","."]
code: 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)
  
imports: import importlib
  import logging
  import os
  import uuid
  from copy import deepcopy
  from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union, cast
  import torch
  from lightning_utilities.core.imports import RequirementCache
  from typing_extensions import override
  import lightning.pytorch
  from lightning.fabric.utilities.types import _TORCH_LRSCHEDULER
  from lightning.pytorch.callbacks import Callback
  from lightning.pytorch.utilities.exceptions import MisconfigurationException
  from lightning.pytorch.utilities.parsing import lightning_hasattr, lightning_setattr
  from lightning.pytorch.utilities.rank_zero import rank_zero_warn
  from lightning.pytorch.utilities.types import STEP_OUTPUT, LRScheduler, LRSchedulerConfig
label_desc: Use image node_id 3 for calling the _LRCallback obj's underlying member method code with example usage: obj.on_train_batch_end(trainer, pl_module, outputs, batch, batch_idx) without return types