code: def test_load_to_same_strategy(self):
      strategy = parameter_server_strategy_v2.ParameterServerStrategyV2(
          self.cluster_resolver
      )
      model_dir = self.get_temp_dir()
      with strategy.scope():
          m = self.Model()
      m.train()
      tf.saved_model.save(m, model_dir)
  
      with strategy.scope():
          loaded = tf.saved_model.load(model_dir)
  
      # Make sure that the variables are created on different devices. SavedModel
      # may load the variables in a different order compared to the creation order
      # so the devices may not be exactly the same as before.
      self.assertTrue(
          (
              "/job:ps/replica:0/task:0" in loaded.v1.device
              and "/job:ps/replica:0/task:1" in loaded.v2.device
          )
          or (
              "/job:ps/replica:0/task:1" in loaded.v1.device
              and "/job:ps/replica:0/task:0" in loaded.v2.device
          )
      )
      self.assertAllEqual(loaded(tf.identity(1)), [6, 6, 6, 6])
  
imports: import os
  from absl.testing import parameterized
  import tensorflow
  import tensorflow.compat.v1
  from tensorflow.python.distribute import combinations
  from tensorflow.python.distribute import multi_worker_test_base
  from tensorflow.python.distribute import parameter_server_strategy_v2
  from tensorflow.python.distribute import sharded_variable
  from tensorflow.python.distribute import strategy_combinations
  from tensorflow.python.distribute import test_util
  from tensorflow.python.distribute import values
  from tensorflow.python.eager import context
  from tensorflow.python.eager import test
  from tensorflow.python.ops import lookup_ops
label_desc: Use image node_id 5 for calling the PSStrategySaveAndLoadTest obj's underlying member method code with example usage: obj.test_load_to_same_strategy() without return types