class_docstr: Abstract class for custom training or validation loops.
  
  The scheduler will only be called at a batch or epoch frequency, depending
  on the user parameter. Be sure to set ``scheduler_step_freq`` in
  ``TorchTrainer`` to either "batch" or "epoch" to increment the scheduler
  correctly during training. If using a learning rate scheduler
  that depends on validation loss, you can use ``trainer.update_scheduler``.
  
  For both training and validation, there are two granularities that
  you can provide customization: per epoch or per batch.
  You do not need to override both.
  
  .. image:: raysgd-custom.jpg
      :scale: 80%
      :align: center
  
  Raises:
      ValueError if multiple models/optimizers/schedulers are provided.
          You are expected to subclass this class if you wish
          to train over multiple models/optimizers/schedulers.
docstr: Computes loss and updates the model over one batch.
  
  This method is responsible for computing the loss and gradient and
  updating the model.
  
  By default, this method implementation assumes that batches
  are in (\*features, labels) format. So we also support multiple inputs
  model. If using amp/fp16 training, it will also scale the loss
  automatically.
  
  You can provide custom loss metrics and training operations if you
  override this method. If overriding this method, you can access model,
  optimizer, criterion via ``self.model``, ``self.optimizer``,
  and ``self.criterion``.
  
  You do not need to override this method if you plan to
  override ``train_epoch``.
  
  Args:
      batch: One item of the validation iterator.
      batch_info (dict): Information dict passed in from ``train_epoch``.
  
  Returns:
      A dictionary of metrics.
          By default, this dictionary contains "loss" and "num_samples".
          "num_samples" corresponds to number of datapoints in the batch.
          However, you can provide any number of other values.
          Consider returning "num_samples" in the metrics because
          by default, ``train_epoch`` uses "num_samples" to
          calculate averages.
docstr_tok: ["Computes","loss","and","updates","the","model","over","one","batch",".","This","method","is","responsible","for","computing","the","loss","and","gradient","and","updating","the","model",".","By","default",",","this","method","implementation","assumes","that","batches","are","in","(","\\","*","features",",","labels",")","format",".","So","we","also","support","multiple","inputs","model",".","If","using","amp\/fp16","training",",","it","will","also","scale","the","loss","automatically",".","You","can","provide","custom","loss","metrics","and","training","operations","if","you","override","this","method",".","If","overriding","this","method",",","you","can","access","model",",","optimizer",",","criterion","via","``","self.model","``",",","``","self.optimizer","``",",","and","``","self.criterion","``",".","You","do","not","need","to","override","this","method","if","you","plan","to","override","``","train_epoch","``",".","Args",":","batch",":","One","item","of","the","validation","iterator",".","batch_info","(","dict",")",":","Information","dict","passed","in","from","``","train_epoch","``",".","Returns",":","A","dictionary","of","metrics",".","By","default",",","this","dictionary","contains","``","loss","''","and","``","num_samples","''",".","``","num_samples","''","corresponds","to","number","of","datapoints","in","the","batch",".","However",",","you","can","provide","any","number","of","other","values",".","Consider","returning","``","num_samples","''","in","the","metrics","because","by","default",",","``","train_epoch","``","uses","``","num_samples","''","to","calculate","averages","."]
code: def train_batch(self, batch, batch_info):
      """Computes loss and updates the model over one batch.
  
      This method is responsible for computing the loss and gradient and
      updating the model.
  
      By default, this method implementation assumes that batches
      are in (\*features, labels) format. So we also support multiple inputs
      model. If using amp/fp16 training, it will also scale the loss
      automatically.
  
      You can provide custom loss metrics and training operations if you
      override this method. If overriding this method, you can access model,
      optimizer, criterion via ``self.model``, ``self.optimizer``,
      and ``self.criterion``.
  
      You do not need to override this method if you plan to
      override ``train_epoch``.
  
      Args:
          batch: One item of the validation iterator.
          batch_info (dict): Information dict passed in from ``train_epoch``.
  
      Returns:
          A dictionary of metrics.
              By default, this dictionary contains "loss" and "num_samples".
              "num_samples" corresponds to number of datapoints in the batch.
              However, you can provide any number of other values.
              Consider returning "num_samples" in the metrics because
              by default, ``train_epoch`` uses "num_samples" to
              calculate averages.
  
      """
      # unpack features into list to support multiple inputs model
      *features, target = batch
      # If features is already a tuple, we don't give it an extra list dimension.
      already_list = isinstance(features[0], tuple) or isinstance(
          features[0], list
      )
      if len(features) == 1 and already_list:
          features = features[0]
  
      # Compute output.
      with self.timers.record("fwd"):
          output = self.model(*features)
          if isinstance(output, tuple) or isinstance(output, list):
              # Then target is also assumed to be a tuple or list.
              loss = self.criterion(*output, *target)
          else:
              loss = self.criterion(output, target)
  
      # Compute gradients in a backward pass.
      with self.timers.record("grad"):
          self.optimizer.zero_grad()
          if self.use_fp16:
              with amp.scale_loss(loss, self.optimizer) as scaled_loss:
                  scaled_loss.backward()
          else:
              loss.backward()
  
      # Call step of optimizer to update model params.
      with self.timers.record("apply"):
          self.optimizer.step()
  
      return {
          "train_loss": loss.item(),
          NUM_SAMPLES: features[0].size(0),
      }
  
imports: import collections
  import torch
  import numpy
  from zoo.orca.learn.metrics import Metric
  from zoo.orca.learn.pytorch.utils import TimerCollection, AverageMeterCollection, NUM_SAMPLES
  from zoo.orca.learn.pytorch.constants import SCHEDULER_STEP_EPOCH, NUM_STEPS, SCHEDULER_STEP_BATCH, SCHEDULER_STEP
  from torch.nn.parallel import DistributedDataParallel
label_desc: Use image node_id 6 for calling the TrainingOperator obj's underlying member method code with example usage: obj.train_batch(batch, batch_info) and returns: dict