code: def test_save_load(self, name, cls):
      n_arms = 5
      b = cls(n_arms=n_arms)
      n_obs_per_arm = torch.randint(0, 100, size=(n_arms,)).float()
      n_success_per_arm = torch.rand(size=(n_arms,)) * n_obs_per_arm
      b.add_batch_observations(
          n_obs_per_arm, n_success_per_arm, n_success_per_arm
      )
  
      avg_rewards_before_save = b.get_avg_reward_values()
  
      if isinstance(b, BaseUCB):
          ucb_scores_before_save = b.get_scores()
  
      f_write = BytesIO()
      torch.save(b, f_write)
      f_write.seek(0)
      f_read = BytesIO(f_write.read())
      f_write.close()
      b_loaded = torch.load(f_read)
      f_read.close()
  
      if isinstance(b, BaseUCB):
          ucb_scores_after_load = b_loaded.get_scores()
          npt.assert_array_equal(
              ucb_scores_before_save.numpy(),
              ucb_scores_after_load.numpy(),
          )  # UCB scores are same before saving and after loading
  
      avg_rewards_after_load = b_loaded.get_avg_reward_values()
      npt.assert_array_equal(
          avg_rewards_before_save.numpy(),
          avg_rewards_after_load.numpy(),
      )  # avg rewards are same before saving and after loading
  
      self.assertListEqual(b.arm_ids, b_loaded.arm_ids)
  
imports: import unittest
  from io import BytesIO
  from itertools import cycle
  from unittest import mock
  import numpy
  import numpy.testing
  import torch
  from parameterized.parameterized import parameterized
  from pytorch_lightning import seed_everything
  from reagent.mab.mab_algorithm import get_arm_indices, place_values_at_indices, randomized_argmax, reindex_multiple_tensors
  from reagent.mab.simulation import BernoilliMAB, compare_bandit_algos, multiple_evaluations_bandit_algo, single_evaluation_bandit_algo
  from reagent.mab.thompson_sampling import BaseThompsonSampling, BernoulliBetaThompson, NormalGammaThompson
  from reagent.mab.ucb import BaseUCB, MetricUCB, UCB1, UCBTuned
label_desc: Use image node_id 5 for calling the TestMAB obj's underlying member method code with example usage: obj.test_save_load(name, cls) without return types