code: def test_call_ucb(self) -> None:
      inputs = []
      # y = x1+x2
      inputs.append(
          CBInput(
              context_arm_features=torch.tensor(
                  [[1.0, 2.0], [1.0, 3.0]]
              ),
              reward=torch.tensor([[3.0], [4.0]]),
          )
      )
      # y = 2x1 + x2
      inputs.append(
          CBInput(
              context_arm_features=torch.tensor(
                  [[2.0, 3.0], [1.0, 5.0]]
              ),
              reward=torch.tensor([[7.0], [7.0]]),
          )
      )
      # y = 2x1 + 2x2
      inputs.append(
          CBInput(
              context_arm_features=torch.tensor(
                  [[0.5, 3.0], [1.8, 5.1]]
              ),
              reward=torch.tensor([[7.0], [13.8]]),
          )
      )
      model = DisjointLinearRegressionUCB(
          num_arms=3, input_dim=2, l2_reg_lambda=0.0
      )
      trainer = DisjointLinUCBTrainer(
          Policy(scorer=model, sampler=GreedyActionSampler())
      )
      trainer.training_step(inputs, batch_idx=0)
      trainer.on_train_epoch_end()
  
      inp = torch.tensor([[1.0, 5.0], [1.0, 6.0]])
      alpha = 1.5
      out = model(inp, ucb_alpha=alpha)
  
      def calculate_mean(x, model_type):
          if model_type == 1:
              return x[0] + x[1]
          elif model_type == 2:
              return 2 * x[0] + x[1]
          elif model_type == 3:
              return 2 * x[0] + 2 * x[1]
  
      expected_out = np.zeros((2, 3))
      for i in range(2):
          x = inp[i]
          for j in range(3):
              expected_out[i, j] = calculate_mean(
                  x, j + 1
              ) + alpha * np.sqrt(
                  x.numpy() @ model.inv_A[j].numpy() @ x.numpy()
              )
      self.assertIsInstance(out, torch.Tensor)
      self.assertEqual(tuple(out.shape), (2, 3))
      npt.assert_allclose(out.numpy(), expected_out, rtol=1e-4)
  
file_path: ReAgent/reagent/test/models/test_disjoint_linucb_predictor.py
imports: import unittest
  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.disjoint_linucb_predictor import batch_quadratic_form_multi_arms, DisjointLinearRegressionUCB
  from reagent.training.cb.disjoint_linucb_trainer import DisjointLinUCBTrainer
label_desc: Use image node_id 1 for calling the TestDisjointLinearRegressionUCB obj's underlying member method code with example usage: obj.test_call_ucb() and returns: unknown, unknown, unknown