code: def test_call(self) -> None:
      bandit = UCB1MAB(ucb_alpha=1.0)
      policy = Policy(
          scorer=bandit,
          sampler=GreedyActionSampler(),
      )
      trainer = MABTrainer(policy)
  
      trainer.cb_training_step(self.batch, 0, 0)
  
      pred_arm_ids = torch.tensor([[1, 3], [6, 4], [5, 2]]).long()
      alpha = 1.5
      model_output = bandit(pred_arm_ids, ucb_alpha=alpha)
  
      ucb = model_output["ucb"]
      self.assertIsInstance(ucb, torch.Tensor)
      self.assertEqual(tuple(ucb.shape), (3, 2))
      # check that we return inf for arms which we haven't observed
      self.assertEqual(
          min(ucb.numpy()[[0, 1, 2, 2], [0, 0, 0, 1]]), float("inf")
      )
      # check that we return non-inf for arms which we have observed
      self.assertLess(max(ucb.numpy()[0:1, 1]), float("inf"))
  
      mu = model_output["pred_label"]
      self.assertIsInstance(ucb, torch.Tensor)
      self.assertEqual(tuple(ucb.shape), (3, 2))
      npt.assert_almost_equal(
          mu.numpy(),
          np.array([[0.0, 3.0], [0.0, 4.0], [0.0, 0.0]]),
      )
  
      sigma = model_output["pred_sigma"]
      self.assertIsInstance(sigma, torch.Tensor)
      self.assertEqual(tuple(sigma.shape), (3, 2))
  
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 2 for calling the TestUCB1MAB obj's underlying member method code with example usage: obj.test_call() without return types