Re-enables tf.function in keras_utils' predict_on_batch conversion.
@@ -334,6 +334,7 @@ class _KerasModel(model_lib.Model):
def input_spec(self):
return self._input_spec
+ @tf.function
def predict_on_batch(self, x, training=True):
return self._keras_model(x, training=training)
