code: def test_model_with_loaded_layer(self, strategy):
      # When a model is loaded under strategy, we wrap it so that when it's passed
      # to strategy.run(), the captured variables resolve to the ones of the
      # current replica. Since the saved tf.function may contain updates to the
      # variables, we don't allow using the model outside of strategy.run().
      #
      # That is to say, a loaded model is different from the original Python one.
      # We need to test save-load-save-load to make sure things work correctly.
  
      class Layer(tf.Module):
          def __init__(self):
              self.v = tf.Variable(1.0)
  
          @tf.function(input_signature=[])
          def __call__(self):
              return self.v.read_value()
  
      class Model(tf.Module):
          def __init__(self, layer):
              self.layer = layer
  
          @tf.function(input_signature=[])
          def __call__(self):
              return self.layer()
  
      layer_export_dir = self.get_temp_dir()
      tf.saved_model.save(Layer(), layer_export_dir)
  
      with strategy.scope():
          m = Model(tf.saved_model.load(layer_export_dir))
          export_dir = self.get_temp_dir()
          # Saving a ConcreteFunction should raise an error.
          with self.assertRaisesRegex(
              ValueError,
              "saving a tf.function with input_signature instead",
          ):
              tf.saved_model.save(
                  m,
                  export_dir,
                  signatures={
                      "call": m.__call__.get_concrete_function(),
                  },
              )
          tf.saved_model.save(m, export_dir)
  
      loaded = tf.saved_model.load(export_dir)
      self.assertAllEqual(
          self.evaluate(
              strategy.experimental_local_results(strategy.run(loaded))
          ),
          [1.0, 1.0],
      )
  
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 6 for calling the SaveAndLoadForServingTest obj's underlying member method code with example usage: obj.test_model_with_loaded_layer(strategy) and returns: self, self