14b7ac5210cf4d4e4cdffcc7739cc0c9817c6a00,official/vision/beta/serving/semantic_segmentation_test.py,SemanticSegmentationExportTest,_export_from_module,#SemanticSegmentationExportTest#Any#Any#Any#,40

Before Change

    return segmentation_module

  def _export_from_module(self, module, input_type, save_directory):
    if input_type == "image_tensor":
      input_signature = tf.TensorSpec(shape=[None, 112, 112, 3], dtype=tf.uint8)
      signatures = {
          "serving_default":
              module.inference_from_image_tensors.get_concrete_function(
                  input_signature)
      }
    elif input_type == "image_bytes":
      input_signature = tf.TensorSpec(shape=[None], dtype=tf.string)
      signatures = {
          "serving_default":
              module.inference_from_image_bytes.get_concrete_function(
                  input_signature)
      }
    elif input_type == "tf_example":
      input_signature = tf.TensorSpec(shape=[None], dtype=tf.string)
      signatures = {
          "serving_default":
              module.inference_from_tf_example.get_concrete_function(
                  input_signature)
      }
    else:
      raise ValueError("Unrecognized `input_type`")

    tf.saved_model.save(module,
                        save_directory,
                        signatures=signatures)

After Change

    return segmentation_module

  def _export_from_module(self, module, input_type, save_directory):
    signatures = module.get_inference_signatures(
        {input_type: "serving_default"})
    tf.saved_model.save(module, save_directory, signatures=signatures)

  def _get_dummy_input(self, input_type):
    Get dummy input for the given input type.

Instances