def get_saved_model_output_tensors()

in realbook/layers/compatibility.py [0:0]

• 10 lines of code
• 5 McCabe index (conditional complexity)

def get_saved_model_output_tensors(saved_model_or_path: Union[tf.keras.Model, str]) -> List[tf.Tensor]:
    """
    Given a path to a SavedModel or an already loaded SavedModel,
    return a list of its output tensors. Useful for figuring out the
    names of the output tensors for use with SavedModelLayer.
    """
    if isinstance(saved_model_or_path, str):
        savedmodel = tf.saved_model.load(saved_model_or_path)
        model = savedmodel.signatures["serving_default"]
        model._backref = savedmodel  # Without this, the SavedModel will be GC'd too early
    else:
        model = saved_model_or_path
    if hasattr(model, "signatures"):
        model = model.signatures["serving_default"]
    return [tensor for tensor in model.outputs if tensor.dtype != "resource"]