Custom TF model with tf.keras.layers.Embedding

This is already using the new native keras model support :slight_smile: