docstr: Perform a training step on a batch of inputs.
  
  Subclass and override to inject custom behavior.
  
  Args:
      model (:obj:`nn.Module`):
          The model to train.
      inputs (:obj:`Dict[str, Union[torch.Tensor, Any]]`):
          The inputs and targets of the model.
  
          The dictionary will be unpacked before being fed to the model. Most models expect the targets under the
          argument :obj:`labels`. Check your model's documentation for all accepted arguments.
  
  Return:
      :obj:`torch.Tensor`: The tensor with training loss on this batch.
docstr_tok: ["Perform","a","training","step","on","a","batch","of","inputs",".","Subclass","and","override","to","inject","custom","behavior",".","Args",":","model","(",":","obj",":","`","nn.Module","`",")",":","The","model","to","train",".","inputs","(",":","obj",":","`","Dict","[","str",",","Union","[","torch.Tensor",",","Any","]","]","`",")",":","The","inputs","and","targets","of","the","model",".","The","dictionary","will","be","unpacked","before","being","fed","to","the","model",".","Most","models","expect","the","targets","under","the","argument",":","obj",":","`","labels","`",".","Check","your","model","'s","documentation","for","all","accepted","arguments",".","Return",":",":","obj",":","`","torch.Tensor","`",":","The","tensor","with","training","loss","on","this","batch","."]
code: def training_step(
      self,
      model: nn.Module,
      inputs: Dict[str, Union[torch.Tensor, Any]],
  ) -> torch.Tensor:
      """
      Perform a training step on a batch of inputs.
  
      Subclass and override to inject custom behavior.
  
      Args:
          model (:obj:`nn.Module`):
              The model to train.
          inputs (:obj:`Dict[str, Union[torch.Tensor, Any]]`):
              The inputs and targets of the model.
  
              The dictionary will be unpacked before being fed to the model. Most models expect the targets under the
              argument :obj:`labels`. Check your model's documentation for all accepted arguments.
  
      Return:
          :obj:`torch.Tensor`: The tensor with training loss on this batch.
      """
  
      oom = False
      oom_message = ""
      try:
          loss = super().training_step(model, inputs)
          return loss
      except RuntimeError as e:
          if "out of memory" in str(e):
              oom = True
              oom_message = str(e)
              logger.warning(
                  f"ran out of memory {self.oom_batch} on {self.args.local_rank}"
              )
              for k, v in inputs.items():
                  print(k, v.size())
          else:
              raise e
  
      if oom:
          self.oom_batch += 1
          raise RuntimeError(oom_message)
  
imports: import torch
  import torch.nn
  from dataclasses import dataclass, field
  from typing import Union, List, Dict, Tuple, Any, Optional
  from torch.cuda.amp import autocast
  from transformers import Seq2SeqTrainer, Seq2SeqTrainingArguments
  from transformers.trainer_pt_utils import LabelSmoother
  from transformers.trainer import *
  from ..seq2seq.constraint_decoder import get_constraint_decoder
label_desc: Use image node_id 2 for calling the ConstraintSeq2SeqTrainer obj's underlying member method code with example usage: obj.training_step(model, inputs) and returns: loss