1 files changed, 4 insertions, 3 deletions
diff --git a/tensorflow/contrib/lite/testing/tflite_diff_util.cc b/tensorflow/contrib/lite/testing/tflite_diff_util.cc
index f601d3752d..19f34c0a51 100644
--- a/tensorflow/contrib/lite/testing/tflite_diff_util.cc
+++ b/tensorflow/contrib/lite/testing/tflite_diff_util.cc
@@ -25,13 +25,14 @@ limitations under the License.
 namespace tflite {
 namespace testing {
 
-bool RunDiffTest(const DiffOptions& options) {
+bool RunDiffTest(const DiffOptions& options, int num_invocations) {
   std::stringstream tflite_stream;
   if (!GenerateTestSpecFromTensorflowModel(
           tflite_stream, options.tensorflow_model, options.tflite_model,
-          options.input_layer, options.input_layer_type,
-          options.input_layer_shape, options.output_layer))
+          num_invocations, options.input_layer, options.input_layer_type,
+          options.input_layer_shape, options.output_layer)) {
     return false;
+  }
   TfLiteDriver tflite_driver(/*use_nnapi=*/true);
   tflite_driver.LoadModel(options.tflite_model);
   return tflite::testing::ParseAndRunTests(&tflite_stream, &tflite_driver);