TFLite-convert-model.py

class Model(tf.keras.Model):
    ...
    def __call__(self, inputs, training):
        # Input layer
        y = tf.reshape(inputs, self._input_shape)
        y = self.conv1(y)
        y = self.max_pool2d(y)
        y = self.conv2(y)
        y = self.max_pool2d(y)
        y = tf.layers.flatten(y)
        y = self.fc1(y)
        y = self.dropout(y, training=training)
        # Returns a logit layer
        return self.fc2(y)