`tf.keras` is better maintained and has better integration with TensorFlow, so we should move to TensorFlow 2.0 and use `tf.keras`