class_docstr: Manages a PyTorch model for training.
class_docstr_tok: ["Manages","a","PyTorch","model","for","training","."]
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.
  
  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",".","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, callbacks=None):
      """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.
  
      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.
  
      """
      self.batch = batch
      self.call_hook(callbacks=callbacks, fn_name="before_train_iter")
  
      # Compute output.
      with self.timers.record("fwd"):
          self.call_hook(
              callbacks=callbacks, fn_name="on_train_forward"
          )
  
      # Compute gradients in a backward pass.
      with self.timers.record("bwd"):
          self.call_hook(
              callbacks=callbacks, fn_name="on_iter_backward"
          )
  
      loss_item = self.loss.item()
      self.metrics_stats = {
          "train_loss": loss_item,
          NUM_SAMPLES: get_batchsize(batch),
      }
  
      self.global_step += 1
  
      self.call_hook(callbacks=callbacks, fn_name="after_train_iter")
  
      # User should not see batch/loss from last iteration
      if hasattr(self, "batch"):
          del self.batch
  
      if hasattr(self, "output"):
          del self.output
  
      if hasattr(self, "loss"):
          del self.loss
  
imports: from filelock import FileLock
  import logging
  import os
  import copy
  import tempfile
  import math
  import torch
  import torch.nn
  import numpy
  from torch.utils.data import DataLoader, IterableDataset
  from bigdl.orca.learn.metrics import Metric
  from bigdl.orca import OrcaContext
  from bigdl.orca.learn.pytorch import utils
  from bigdl.orca.learn.pytorch.utils import AverageMeterCollection, NUM_SAMPLES, get_batchsize, index_concatenate, split_predict_cols
  from bigdl.orca.learn.pytorch.core import BaseRunner
  from bigdl.dllib.utils.log4Error import invalidInputError
label_desc: Use image node_id 11 for calling the TorchRunner obj's underlying member method code with example usage: obj._train_batch(batch, callbacks) without return types