Fix and re-enable cwise ops benchmark.

PiperOrigin-RevId: 339956677 Change-Id: Id4cea5f63b3f257170c121d74e9cb2995a10b629
2020-10-30 16:05:54 -07:00 · 2020-10-30 16:05:54 -07:00 · 3ad80ca351
commit 3ad80ca351
parent 0b752d9223
1 changed files with 2 additions and 1 deletions
--- a/tensorflow/core/kernels/cwise_ops_test.cc
+++ b/tensorflow/core/kernels/cwise_ops_test.cc
@ -184,7 +184,8 @@ Graph* CubeWithMulSquare(int num) {
  void BM_##DEVICE##_Cube_##Impl(::testing::benchmark::State& state) { \
    const int num = state.range(0);                                    \
                                                                       \
-    test::Benchmark(#DEVICE, Impl(num)).Run(state.iterations());       \
+    test::Benchmark(#DEVICE, Impl(num), /*old_benchmark_api*/ false)   \
+        .Run(state);                                                   \
    const int64 tot = static_cast<int64>(state.iterations()) * num;    \
    state.SetItemsProcessed(tot);                                      \
    state.SetBytesProcessed(tot * sizeof(float));                      \