code: def test_process_all_data(self):
      state_dict_before = copy.deepcopy(self.eval_module.state_dict())
      self.eval_module._process_all_data(self.batch)
      state_dict_after = copy.deepcopy(self.eval_module.state_dict())
  
      # sum_weight_all_data_local got updated properly
      self.assertAlmostEqual(
          state_dict_after["sum_weight_all_data_local"].item()
          - state_dict_before["sum_weight_all_data_local"].item(),
          len(self.batch),
      )
      # sum_weight_all_data didn't change (bcs we haven't aggregated across instances yet)
      self.assertAlmostEqual(
          state_dict_after["sum_weight_all_data"].item(),
          state_dict_before["sum_weight_all_data"].item(),
      )
  
      # sum_reward_weighted_accepted_local got updated properly
      self.assertAlmostEqual(
          state_dict_after["sum_reward_weighted_accepted_local"].item(),
          state_dict_before[
              "sum_reward_weighted_accepted_local"
          ].item(),
      )
      # sum_reward_weighted_accepted didn't change (bcs we haven't aggregated across instances yet)
      self.assertAlmostEqual(
          state_dict_after["sum_reward_weighted_accepted"].item(),
          state_dict_before["sum_reward_weighted_accepted"].item(),
      )
  
      # sum_weight and sum_reward_weighted didn't change (as well as local values)
      self.assertAlmostEqual(
          state_dict_after["sum_weight_accepted"].item(),
          state_dict_before["sum_weight_accepted"].item(),
      )
      self.assertAlmostEqual(
          state_dict_after["sum_weight_accepted_local"].item(),
          state_dict_before["sum_weight_accepted_local"].item(),
      )
      self.assertAlmostEqual(
          state_dict_after["sum_reward_weighted_accepted"].item(),
          state_dict_before["sum_reward_weighted_accepted"].item(),
      )
      self.assertAlmostEqual(
          state_dict_after["sum_reward_weighted_accepted_local"].item(),
          state_dict_before[
              "sum_reward_weighted_accepted_local"
          ].item(),
      )
  
imports: import copy
  import unittest
  from dataclasses import replace
  from unittest.mock import MagicMock
  import torch
  from pytorch_lightning.loggers import TensorBoardLogger
  from reagent.core.types import CBInput
  from reagent.evaluation.cb.policy_evaluator import PolicyEvaluator
  from reagent.evaluation.cb.utils import add_importance_weights
  from reagent.models.linear_regression import LinearRegressionUCB
label_desc: Use image node_id 2 for calling the TestPolicyEvaluator obj's underlying member method code with example usage: obj.test_process_all_data() without return types