'How to retrieve output of exported function?

When calling a function from an exported TensorFlow model, I recieve two strings ("output_0", "output_1") instead of the actual model. How could I retrueve the Tensor associated with this string to access the output?

exported model:

class OneStep(tf.keras.Model):
  ...
  @tf.function(input_signature=[tf.TensorSpec(shape=[None], dtype=tf.string), tf.TensorSpec(shape=[None, None], dtype=tf.float32)])
  def generate_one_step(self, inputs, states):

  ...
  @tf.function(input_signature=[tf.TensorSpec(shape=[None], dtype=tf.string)])
  def generate_one_step_none(self, inputs):

...
tf.saved_model.save(one_step_model, 'one_Step', signatures={ "generate_one_step": one_step_model.generate_one_step, "generate_one_step_none": one_step_model.generate_one_step_none})

code to import:

one_step = tf.saved_model.load('one_step')
step_gen = one_step.signatures["generate_one_step"]
step_gen_none = one_step.signatures["generate_one_step_none"]

next_char = tf.constant(['Test'], tf.string)

a, b = step_gen_none(inputs=next_char)
print(a,b) # returns "input_0", "input_1"


Sources

This article follows the attribution requirements of Stack Overflow and is licensed under CC BY-SA 3.0.

Source: Stack Overflow

Solution Source