Model Class with AutoGraph

Model_with_AutoGraph.py

class MyModel(tf.keras.Model):
  def __init__(self, keep_probability=0.2):
    super(MyModel, self).__init__()
    self.dense1 = tf.keras.layers.Dense(4)
    self.dense2 = tf.keras.layers.Dense(5)
    self.keep_probability = keep_probability

  @tf.function
  def call(self, inputs, training=True):
    y = self.dense2(self.dense1(inputs))
    if training:
      return tf.nn.dropout(y, self.keep_probability)
    else:
      return y

model = MyModel()
model(x, training=True)  # executes a graph, with dropout
model(x, training=False) # executes a graph, without dropout