class_docstr: Mixin class for applying the R3F method, to apply R3F with any model
  inherit the class and implement the abstract functions.
  
  For more details: https://arxiv.org/abs/2008.03156
class_docstr_tok: ["Mixin","class","for","applying","the","R3F","method",",","to","apply","R3F","with","any","model","inherit","the","class","and","implement","the","abstract","functions",".","For","more","details",":","https",":","\/\/arxiv.org\/abs\/2008.03156"]
docstr: Runs training over a batch with the R3F method, training will use R3F
  while eval and test do not.
docstr_tok: ["Runs","training","over","a","batch","with","the","R3F","method",",","training","will","use","R3F","while","eval","and","test","do","not","."]
code: def train_batch(cls, model, batch, state=None):
      """
      Runs training over a batch with the R3F method, training will use R3F
      while eval and test do not.
      """
  
      # Forward pass through the network.
      model_inputs = model.arrange_model_inputs(batch)
      model_context = model.arrange_model_context(batch)
      targets = model.arrange_targets(batch)
  
      sample_size = model.get_sample_size(
          model_inputs=model_inputs, targets=targets
      )
  
      # get embedding
      r3f_loss_term = torch.tensor(0)
      if state and state.stage == Stage.TRAIN:
          # during training run R3F forward calls
          model_outputs, noise_model_outputs = model(
              *model_inputs, use_r3f=True
          )
  
          r3f_loss_term = model.get_r3f_loss_terms(
              model_outputs,
              noise_model_outputs,
              sample_size=sample_size,
          )
      else:
          # during eval and test don't run R3F forward
          model_outputs = model(*model_inputs, use_r3f=False)
  
      # Add stage to context.
      if state:
          if model_context is None:
              model_context = {
                  "stage": state.stage,
                  "epoch": state.epoch,
              }
          else:
              model_context["stage"] = state.stage
              model_context["epoch"] = state.epoch
  
      # Compute loss and predictions.
      loss = maybe_float(
          model.get_loss(model_outputs, targets, model_context)
      )
  
      # add R3F loss term
      loss = loss + r3f_loss_term.to(loss.device)
  
      predictions, scores = model.get_pred(
          model_outputs, context=model_context
      )
  
      # Pack results and return them.
      metric_data = (predictions, targets, scores, loss, model_inputs)
      return loss, metric_data
  
imports: from contextlib import AbstractContextManager
  from enum import Enum
  from typing import Dict
  import torch
  import torch.nn.functional
  from pytext.common.constants import Stage
  from pytext.config import ConfigBase
  from pytext.utils.precision import maybe_float
label_desc: Use image node_id 9 for calling the R3FPyTextMixin obj's underlying member method code with example usage: obj.train_batch(cls, model, batch, state) and returns: loss,  metric_data