From d0eeef269d1f18961c6cd0a8d80ede564626fbc3 Mon Sep 17 00:00:00 2001 From: "A. Unique TensorFlower" Date: Wed, 24 Jul 2019 11:36:35 -0700 Subject: [PATCH] Use tf.function's default autograph=True in saved_model/integration_test. There is currently no demonstrable need to do something non-obvious here. PiperOrigin-RevId: 259782705 --- .../examples/saved_model/integration_tests/export_mnist_cnn.py | 2 +- .../examples/saved_model/integration_tests/export_rnn_cell.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/tensorflow/examples/saved_model/integration_tests/export_mnist_cnn.py b/tensorflow/examples/saved_model/integration_tests/export_mnist_cnn.py index 74981b5fbf7..6b94fda0f34 100644 --- a/tensorflow/examples/saved_model/integration_tests/export_mnist_cnn.py +++ b/tensorflow/examples/saved_model/integration_tests/export_mnist_cnn.py @@ -117,7 +117,7 @@ def wrap_keras_model_for_export(model, batch_input_shape, # the desired argspec. def wrapped(*args, **kwargs): # TODO(arnoegw): Can we use call_fn itself? return call_fn(*args, **kwargs) - traced_call_fn = tf.function(autograph=False)( + traced_call_fn = tf.function( tf_decorator.make_decorator(call_fn, wrapped, decorator_argspec=argspec)) # Now we need to trigger traces for all supported combinations of the diff --git a/tensorflow/examples/saved_model/integration_tests/export_rnn_cell.py b/tensorflow/examples/saved_model/integration_tests/export_rnn_cell.py index 876e3004bca..6a2853f0617 100644 --- a/tensorflow/examples/saved_model/integration_tests/export_rnn_cell.py +++ b/tensorflow/examples/saved_model/integration_tests/export_rnn_cell.py @@ -37,7 +37,7 @@ def main(argv): root.rnn_cell = tf.keras.layers.LSTMCell(units=10, recurrent_initializer=None) # Wrap the rnn_cell.__call__ function and assign to next_state. - root.next_state = tf.function(root.rnn_cell.__call__, autograph=False) + root.next_state = tf.function(root.rnn_cell.__call__) # Wrap the rnn_cell.get_initial_function using a decorator and assign to an # attribute with the same name.