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: Check that the input batch satisfies the following assumptions:
      1. context_arm_features is 3D with dimensions: batch_size, arm_count, feature_dim
      2. Label and action are not none
      3. Batch size is the same for action, reward and context_arm_features
docstr_tok: ["Check","that","the","input","batch","satisfies","the","following","assumptions",":","1.","context_arm_features","is","3D","with","dimensions",":","batch_size",",","arm_count",",","feature_dim","2",".","Label","and","action","are","not","none","3",".","Batch","size","is","the","same","for","action",",","reward","and","context_arm_features"]
code: def _check_input(
      self, batch: CBInput, offline_eval: bool = False
  ) -> None:
      """
      Check that the input batch satisfies the following assumptions:
          1. context_arm_features is 3D with dimensions: batch_size, arm_count, feature_dim
          2. Label and action are not none
          3. Batch size is the same for action, reward and context_arm_features
      """
      assert batch.context_arm_features.ndim == 3
      assert batch.label is not None
      assert batch.action is not None
      assert len(batch.action) == len(batch.label)
      assert len(batch.action) == batch.context_arm_features.shape[0]
      if offline_eval:
          assert batch.reward is not None
          assert len(batch.action) == len(batch.reward)
  
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 2 for calling the BaseCBTrainerWithEval obj's underlying member method code with example usage: obj._check_input(batch, offline_eval) without return types