class_docstr: A training engine hybrid pipeline, data, and model parallel training.
  
  This engine is created by ``deepspeed.initialize()`` when a :class:`PipelineModule`
  is provided.
docstr: Progress the pipeline to train the next batch of data. The engine will ingest
  ``self.train_batch_size()`` total samples collectively across all workers.
  
  
  An iterator that over training data should be provided as an argument
  unless ``deepspeed.initialize()`` was provided a training set. In that event,
  the training data will automatically be read.
  
  
  .. warning::
      A total of ``self.gradient_accumulation_steps()`` entries will be pulled
      from ``data_iter`` by each pipeline. There must be sufficient
      data left in ``data_iter`` or else a ``StopIteration`` will halt training.
  
      DeepSpeed provides a convenience class :class:`deepspeed.utils.RepeatingLoader`
      that wraps data loaders to automatically restart upon a ``StopIteration``.
  
  Args:
      data_iter (Iterator, optional): Iterator of training data.
  
  Returns:
      The arithmetic mean of the losses computed this batch.
docstr_tok: ["Progress","the","pipeline","to","train","the","next","batch","of","data",".","The","engine","will","ingest","``","self.train_batch_size","(",")","``","total","samples","collectively","across","all","workers",".","An","iterator","that","over","training","data","should","be","provided","as","an","argument","unless","``","deepspeed.initialize","(",")","``","was","provided","a","training","set",".","In","that","event",",","the","training","data","will","automatically","be","read",".","..","warning",":",":","A","total","of","``","self.gradient_accumulation_steps","(",")","``","entries","will","be","pulled","from","``","data_iter","``","by","each","pipeline",".","There","must","be","sufficient","data","left","in","``","data_iter","``","or","else","a","``","StopIteration","``","will","halt","training",".","DeepSpeed","provides","a","convenience","class",":","class",":","`","deepspeed.utils.RepeatingLoader","`","that","wraps","data","loaders","to","automatically","restart","upon","a","``","StopIteration","``",".","Args",":","data_iter","(","Iterator",",","optional",")",":","Iterator","of","training","data",".","Returns",":","The","arithmetic","mean","of","the","losses","computed","this","batch","."]
code: def train_batch(self, data_iter=None):
      """Progress the pipeline to train the next batch of data. The engine will ingest
      ``self.train_batch_size()`` total samples collectively across all workers.
  
  
      An iterator that over training data should be provided as an argument
      unless ``deepspeed.initialize()`` was provided a training set. In that event,
      the training data will automatically be read.
  
  
      .. warning::
          A total of ``self.gradient_accumulation_steps()`` entries will be pulled
          from ``data_iter`` by each pipeline. There must be sufficient
          data left in ``data_iter`` or else a ``StopIteration`` will halt training.
  
          DeepSpeed provides a convenience class :class:`deepspeed.utils.RepeatingLoader`
          that wraps data loaders to automatically restart upon a ``StopIteration``.
  
      Args:
          data_iter (Iterator, optional): Iterator of training data.
  
      Returns:
          The arithmetic mean of the losses computed this batch.
      """
      if not torch._C.is_grad_enabled():
          raise RuntimeError(
              f"train_batch() requires gradients enabled. Use eval_batch() instead."
          )
  
      # Curriculum learning could change activation shape
      if self.curriculum_enabled():
          new_difficulty = self.curriculum_scheduler.update_difficulty(
              self.global_steps + 1
          )
          if (
              self.global_steps == 0
              or self.curriculum_scheduler.first_step
          ):
              self.reset_activation_shape()
              self.curriculum_scheduler.first_step = False
          elif (
              new_difficulty
              != self.curriculum_scheduler.get_difficulty(
                  self.global_steps
              )
          ):
              self.reset_activation_shape()
  
      if data_iter:
          self.set_dataiterator(data_iter)
  
      self.module.train()
      self.total_loss = None
      self._compute_loss = True
  
      # Do the work
      self.timers("train_batch").start()
      sched = schedule.TrainSchedule(
          micro_batches=self.micro_batches,
          stages=self.num_stages,
          stage_id=self.stage_id,
      )
      self._exec_schedule(sched)
      self.agg_train_loss = self._aggregate_total_loss()
  
      self.timers("train_batch").stop()
  
      if self.global_steps % self.steps_per_print() == 0:
          if self.global_rank == 0:
              elapsed = (
                  self.timers("train_batch").elapsed(reset=True)
                  / 1000.0
              )
              iter_time = elapsed / self.steps_per_print()
              tput = self.train_batch_size() / iter_time
              print(
                  f"steps: {self.global_steps} "
                  f"loss: {self.agg_train_loss:0.4f} "
                  f"iter time (s): {iter_time:0.3f} "
                  f"samples/sec: {tput:0.3f}"
              )
  
      # Monitoring
      if self.global_rank == 0 and self.monitor.enabled:
          self.summary_events = [
              (
                  f"Train/Samples/train_loss",
                  self.agg_train_loss.mean().item(),
                  self.global_samples,
              )
          ]
          self.monitor.write_events(self.summary_events)
  
      if (
          self.wall_clock_breakdown()
          and self.global_steps % self.steps_per_print() == 0
      ):
          self.timers.log(
              [
                  "pipe_send_output",
                  "pipe_send_grad",
                  "pipe_recv_input",
                  "pipe_recv_grad",
              ]
          )
  
      # TODO: should return precisely what loss returned and allow others to be queried?
      return self.agg_train_loss
  
imports: from types import MethodType
  import torch
  from deepspeed import comm
  from deepspeed.utils import logger
  from deepspeed.utils.timer import ThroughputTimer
  from ..engine import DeepSpeedEngine, MEMORY_OPT_ALLREDUCE_SIZE
  from ..utils import PartitionedTensor
  from ..dataloader import RepeatingLoader
  from .module import PipelineModule, PipelineError
  from .None import p2p
  from .None import schedule
label_desc: Use image node_id 9 for calling the PipelineEngine obj's underlying member method code with example usage: obj.train_batch(data_iter) and returns: self