class_docstr: Monitor the speed of each step and each epoch.
      
code: def on_train_batch_start(
      self,
      trainer: "pl.Trainer",
      pl_module: "pl.LightningModule",
      batch: Any,
      batch_idx: int,
  ) -> None:
      if self._log_stats.intra_step_time:
          self._snap_intra_step_time = time.time()
  
      if not trainer._logger_connector.should_update_logs:
          return
  
      logs = {}
      if self._log_stats.inter_step_time and self._snap_inter_step_time:
          # First log at beginning of second step
          logs["time/inter_step (ms)"] = (
              time.time() - self._snap_inter_step_time
          ) * 1000
  
      if trainer.logger is not None:
          trainer.logger.log_metrics(logs, step=trainer.global_step)
  
file_path: flash-attention/training/src/callbacks/speed_monitor.py
imports: import time
  from typing import Any
  from pytorch_lightning import Callback, Trainer
  from pytorch_lightning.utilities import rank_zero_only
  from pytorch_lightning.utilities.parsing import AttributeDict
  from pytorch_lightning.utilities.types import STEP_OUTPUT
label_desc: Use image node_id 6 for calling the SpeedMonitor obj's underlying member method code with example usage: obj.on_train_batch_start(trainer, pl_module, batch, batch_idx) without return types