class_docstr: The base class for Contextual Bandit models. A minimal implementation of a specific model requires providing a specific
      implementation for only the cb_training_step() method.
  The main functionality implemented in this base class is the integration of Offline Evaluation into the training loop.
class_docstr_tok: ["The","base","class","for","Contextual","Bandit","models",".","A","minimal","implementation","of","a","specific","model","requires","providing","a","specific","implementation","for","only","the","cb_training_step","(",")","method",".","The","main","functionality","implemented","in","this","base","class","is","the","integration","of","Offline","Evaluation","into","the","training","loop","."]
docstr: This method combines 2 things in order to enable Offline Evaluation of non-stationary CB algorithms:
  1. If offline evaluator is defined, it will pre-process the batch - keep track of the reward and filter out some observations.
  2. The filtered batch will be fed to the cb_training_step() method, which implements the actual training logic.
  
  DO NOT OVERRIDE THIS METHOD IN SUBCLASSES, IT'S @final. Instead, override cb_training_step().
docstr_tok: ["This","method","combines","2","things","in","order","to","enable","Offline","Evaluation","of","non-stationary","CB","algorithms",":","1",".","If","offline","evaluator","is","defined",",","it","will","pre-process","the","batch","-","keep","track","of","the","reward","and","filter","out","some","observations",".","2",".","The","filtered","batch","will","be","fed","to","the","cb_training_step","(",")","method",",","which","implements","the","actual","training","logic",".","DO","NOT","OVERRIDE","THIS","METHOD","IN","SUBCLASSES",",","IT","'S","@","final",".","Instead",",","override","cb_training_step","(",")","."]
code: def training_step(
      self, batch: CBInput, batch_idx: int, optimizer_idx: int = 0
  ) -> Optional[torch.Tensor]:
      """
      This method combines 2 things in order to enable Offline Evaluation of non-stationary CB algorithms:
      1. If offline evaluator is defined, it will pre-process the batch - keep track of the reward and filter out some observations.
      2. The filtered batch will be fed to the cb_training_step() method, which implements the actual training logic.
  
      DO NOT OVERRIDE THIS METHOD IN SUBCLASSES, IT'S @final. Instead, override cb_training_step().
      """
      eval_module = (
          self.eval_module
      )  # assign to local var to keep pyre happy
      offline_eval_enabled = eval_module is not None
      self._check_input(batch, offline_eval=offline_eval_enabled)
      if offline_eval_enabled:
          assert eval_module is not None
          # update the model if we've processed enough samples
          eval_model_update_critical_weight = (
              self.eval_model_update_critical_weight
          )
          if eval_model_update_critical_weight is not None:
              if (
                  eval_module.sum_weight_since_update_local.item()
                  >= eval_model_update_critical_weight
              ):
                  logger.info(
                      f"Updating the evaluated model after {eval_module.sum_weight_since_update_local.item()} observations"
                  )
                  eval_module.update_eval_model(self.scorer)
                  eval_module.sum_weight_since_update_local.zero_()
                  eval_module.num_eval_model_updates += 1
                  eval_module._aggregate_across_instances()
                  eval_module.log_metrics(step=self.global_step)
          with torch.no_grad():
              # TODO: unify contextual and non-contextual logic by passing
              # CBInput object to the model and letting the model pick which element(s) to use
              if isinstance(eval_module.eval_model, UCBBaseModel):
                  model_output = eval_module.eval_model(
                      batch.context_arm_features
                  )
              elif isinstance(eval_module.eval_model, MABBaseModel):
                  model_output = eval_module.eval_model(batch.arms)
              ucb = model_output["ucb"]
              eval_scores = ucb
              model_actions = get_model_actions(
                  eval_scores, batch.arm_presence
              )
          new_batch = eval_module.ingest_batch(batch, model_actions)
          eval_module.sum_weight_since_update_local += (
              batch.weight.sum()
              if batch.weight is not None
              else len(batch)
          )
      else:
          new_batch = batch
      new_batch = add_chosen_arm_features(new_batch)
      ret = self.cb_training_step(
          new_batch, batch_idx, optimizer_idx
      )  # pyre-ignore
      if isinstance(new_batch, CBInput):
          # recmetrics currently supported only for joint model
          if isinstance(self.scorer, UCBBaseModel):
              assert new_batch.features_of_chosen_arm is not None
              x = new_batch.features_of_chosen_arm
          elif isinstance(self.scorer, MABBaseModel):
              assert new_batch.chosen_arm_id is not None
              x = new_batch.chosen_arm_id
          else:
              raise RuntimeError(
                  f"Scorer type {type(self.scorer)} not supported"
              )
          self._update_recmetrics(new_batch, batch_idx, x)
      return ret
  
imports: import logging
  from abc import ABC, abstractmethod
  from typing import final, Optional
  import torch
  from reagent.core.types import CBInput
  from reagent.core.utils import get_rank
  from reagent.evaluation.cb.base_evaluator import BaseOfflineEval
  from reagent.models.cb_base_model import UCBBaseModel
  from reagent.models.mab import MABBaseModel
  from reagent.training.cb.utils import add_chosen_arm_features, get_model_actions
  from reagent.training.reagent_lightning_module import ReAgentLightningModule
  from torchrec.metrics.metric_module import RecMetricModule
label_desc: Use image node_id 5 for calling the BaseCBTrainerWithEval obj's underlying member method code with example usage: obj.training_step(batch, batch_idx, optimizer_idx) and returns: ret