tf.keras.layers not supported

[cweill]

We consistently see issues when users use tf.keras.layers, which are fixed when they switch to the tf.layers equivalent.

Related to #1.

[chamorajg]

Can I work on this issue @cweill ?