code: def test_add_chosen_arm_features(self):
      all_arms_features = torch.tensor(
          [[[1.0, 2.0], [3.0, 4.0]], [[5.0, 6.0], [7.0, 8.0]]],
          dtype=torch.float,
      )
      arm_ids = torch.tensor([[10, 13], [15, 20]], dtype=torch.long)
      actions = torch.tensor([[1], [0]], dtype=torch.long)
      batch = CBInput(
          context_arm_features=all_arms_features,
          action=actions,
          arms=arm_ids,
      )
      new_batch = add_chosen_arm_features(batch)
      npt.assert_equal(
          new_batch.features_of_chosen_arm.numpy(),
          np.array([[3.0, 4.0], [5.0, 6.0]]),
      )
      npt.assert_equal(
          new_batch.chosen_arm_id.numpy(), np.array([[13], [15]])
      )
  
file_path: ReAgent/reagent/test/training/cb/test_utils.py
imports: import unittest
  import numpy
  import torch
  from numpy import testing
  from reagent.core.types import CBInput
  from reagent.training.cb.utils import add_chosen_arm_features, argmax_random_tie_breaks, get_model_actions
label_desc: Use image node_id 1 for calling the TestCButils obj's underlying member method code with example usage: obj.test_add_chosen_arm_features() without return types