code: def test_call_ucb(self) -> None:
      x = torch.tensor([[1.0, 2.0], [1.0, 3.0]])  # y=x+1
      y = torch.tensor([3.0, 4.0])
      model = LinearRegressionUCB(2, l2_reg_lambda=0.0)
      trainer = LinUCBTrainer(
          Policy(scorer=model, sampler=GreedyActionSampler())
      )
      trainer.update_params(x, y)
  
      inp = torch.tensor([[1.0, 5.0], [1.0, 6.0]])
      alpha = 1.5
      model_output = model(inp, ucb_alpha=alpha)
      ucb = model_output["ucb"]
  
      expected_out = np.zeros(2)
      expected_out[0] = 6.0 + alpha * np.sqrt(
          inp[0].numpy()
          @ (model.inv_avg_A / model.sum_weight).numpy()
          @ inp[0].numpy()
      )
      expected_out[1] = 7.0 + alpha * np.sqrt(
          inp[1].numpy()
          @ (model.inv_avg_A / model.sum_weight).numpy()
          @ inp[1].numpy()
      )
  
      self.assertIsInstance(ucb, torch.Tensor)
      self.assertEqual(tuple(ucb.shape), (2,))
      npt.assert_allclose(ucb.numpy(), expected_out, rtol=1e-4)
  
file_path: ReAgent/reagent/test/models/test_linear_regression_ucb.py
imports: import unittest
  import numpy
  import numpy.testing
  import torch
  from reagent.gym.policies.policy import Policy
  from reagent.gym.policies.samplers.discrete_sampler import GreedyActionSampler
  from reagent.models.linear_regression import batch_quadratic_form, LinearRegressionUCB
  from reagent.training.cb.linucb_trainer import LinUCBTrainer
label_desc: Use image node_id 2 for calling the TestLinearRegressionUCB obj's underlying member method code with example usage: obj.test_call_ucb() without return types