title: Tensorflow 2.0 how to export predictions and training-mode from Keras model
body: <p>I have a model sub-classed from tf.keras.model . I have write a call -and predict methods. when I export the model it seems that only the output from call method is serialized. The simple code to illustrate the problem is below</p>
  
  <pre><code>class SimpleModel(tf.keras.Model):
  
  def __init__(self):
      super(SimpleModel, self).__init__()
  
      self.layer1 = keras.layers.Flatten(input_shape=(28, 28))
      self.layer2 = keras.layers.Dense(128, activation='relu')
      self.dropout = keras.layers.Dropout(0.5)
      self.layer3 = keras.layers.Dense(10, activation='softmax')
  
  def call(self, x, training=False):
  
      x = self.layer1(x)
      x = self.layer2(x)
      if training:
          x = self.dropout(x)
      return self.layer3(x)
  
  def predict(self, x):
  
      return tf.argmax(self(x))
  
  fashion_mnist = keras.datasets.fashion_mnist
  
  (train_images, train_labels), (test_images, test_labels) = fashion_mnist.load_data()
  
  class_names = ['T-shirt/top', 'Trouser', 'Pullover', 'Dress', 'Coat',
             'Sandal', 'Shirt', 'Sneaker', 'Bag', 'Ankle boot']
  
  model = SimpleModel()
  
  model.compile(optimizer='adam',
            loss='sparse_categorical_crossentropy',
            metrics=['accuracy'])
  
  model.fit(train_images, train_labels, epochs=10)
  
  model.save('tf_test', save_format='tf')
  </code></pre>
  
  <p>When I inspect the saved model using the saved-model-cli the output is as below</p>
  
  <pre><code>The given SavedModel SignatureDef contains the following input(s):
    inputs['input_1'] tensor_info:
        dtype: DT_UINT8
        shape: (-1, 28, 28)
        name: serving_default_input_1:0
  The given SavedModel SignatureDef contains the following output(s):
    outputs['output_1'] tensor_info:
        dtype: DT_FLOAT
        shape: (-1, 10)
        name: StatefulPartitionedCall:0
  Method name is: tensorflow/serving/predict
  </code></pre>
  
  <p>The output is exactly same without the predict method. So how can I get the prediction tensor included in the serialization and does the training mode need to be serialized. In tensorflow 1.x I can simply save the prediction tensor and training mode tensor using the following line of code</p>
  
  <pre><code>tf.saved_model.simple_save(sess, 'tf_test', inputs={"x": x, "mode": training_mode}, outputs={"predictions": predictions})
  </code></pre>