code: def on_train_step_end(
      self,
      steps_completed: int,
      metrics: Dict[str, Any],
      batch_metrics: Optional[List[Dict[str, Any]]] = None,
  ) -> None:
      logger.debug("Write training metrics for TensorBoard")
      metrics_seen = set()
  
      # Log all batch metrics.
      if batch_metrics:
          for batch_idx, batch in enumerate(batch_metrics):
              batches_seen = (
                  steps_completed - len(batch_metrics) + batch_idx
              )
              for name, value in batch.items():
                  self._maybe_write_metric(name, value, batches_seen)
                  metrics_seen.add(name)
  
      # Log avg metrics which were calculated by a custom reducer and are not in batch metrics.
      for name, value in metrics.items():
          if name in metrics_seen:
              continue
          self._maybe_write_metric(name, value, steps_completed)
  
      self.writer.reset()
  
file_path: determined/harness/determined/tensorboard/metric_writers/callback.py
imports: import abc
  import logging
  from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union
  from determined import util
label_desc: Use image node_id 3 for calling the BatchMetricWriter obj's underlying member method code with example usage: obj.on_train_step_end(steps_completed, metrics, batch_metrics) without return types