code: def test_add_importance_weights(self):
      model_actions = torch.tensor([[1], [1]], dtype=torch.long)
      new_batch = add_importance_weights(self.batch, model_actions)
      # everything except weights should remain the same in the new batch
      for name in ["context_arm_features", "action", "reward"]:
          npt.assert_allclose(
              getattr(self.batch, name).numpy(),
              getattr(new_batch, name).numpy(),
          )
  
      # data weights should be unchanged
      npt.assert_allclose(
          new_batch.weight.numpy(), self.batch.weight.numpy()
      )
      # weights should be zero-ed out where action!= model_action
      self.assertEqual(new_batch.effective_weight[0, 0].item(), 0.0)
      # weights should be multiplied inverse probability of logged action where action == model_action
      self.assertEqual(
          new_batch.effective_weight[1, 0].item(),
          self.batch.weight[1, 0].item() * 2,
      )
  
file_path: ReAgent/reagent/test/evaluation/cb/test_utils.py
imports: import unittest
  import numpy.testing
  import torch
  from reagent.core.types import CBInput
  from reagent.evaluation.cb.utils import add_importance_weights
label_desc: Use image node_id 2 for calling the TestCBEvalUtils obj's underlying member method code with example usage: obj.test_add_importance_weights() without return types