code: def setUp(self) -> None:
      super().setUp()
      self.batch = CBInput.from_dict(
          {
              "context_arm_features": torch.zeros(2, 3, 1),
              "action": torch.Tensor([[2], [0]]).long(),
              "reward": torch.Tensor([[3.0], [4.0]]),
              "arms": torch.Tensor([[1, 2, 3], [4, 5, 6]]).long(),
          }
      )
  
file_path: ReAgent/reagent/test/models/test_mab.py
imports: import unittest
  from dataclasses import replace
  import numpy
  import numpy.testing
  import torch
  from reagent.core.types import CBInput
  from reagent.gym.policies.policy import Policy
  from reagent.gym.policies.samplers.discrete_sampler import GreedyActionSampler
  from reagent.models.mab import UCB1MAB
  from reagent.training.cb.mab_trainer import MABTrainer
label_desc: Use image node_id 1 for calling the TestUCB1MAB obj's underlying member method code with example usage: obj.setUp() without return types