diff options
Diffstat (limited to 'tensorflow/tools/benchmark/benchmark_model_test.cc')
| -rw-r--r-- | tensorflow/tools/benchmark/benchmark_model_test.cc | 5 |
1 files changed, 4 insertions, 1 deletions
diff --git a/tensorflow/tools/benchmark/benchmark_model_test.cc b/tensorflow/tools/benchmark/benchmark_model_test.cc index ee7f24c0cf..c8deeb3625 100644 --- a/tensorflow/tools/benchmark/benchmark_model_test.cc +++ b/tensorflow/tools/benchmark/benchmark_model_test.cc @@ -62,8 +62,11 @@ TEST(BenchmarkModelTest, InitializeAndRun) { std::unique_ptr<StatSummarizer> stats; stats.reset(new tensorflow::StatSummarizer(*(loaded_graph_def.get()))); int64 time; + int64 num_runs; TF_ASSERT_OK(benchmark_model::TimeMultipleRuns( - 0.0, 10, {input}, {output_name}, session.get(), stats.get(), &time)); + 0.0, 10, 0.0, {input}, {output_name}, session.get(), stats.get(), &time, + &num_runs)); + ASSERT_EQ(num_runs, 10); } } // namespace |