Make BERT SQuAD pure eager test run shorter to avoid timeout

PiperOrigin-RevId: 298734328

Make BERT SQuAD pure eager test run shorter to avoid timeout
PiperOrigin-RevId: 298734328
75d13042 · Zongwei Zhou · A. Unique TensorFlower · b0ba8502 · 75d13042
Commit 75d13042 authored Mar 03, 2020 by Zongwei Zhou Committed by A. Unique TensorFlower Mar 03, 2020
Show whitespace changes
Inline Side-by-side

Showing with 1 addition and 1 deletion

official/benchmark/bert_squad_benchmark.py official/benchmark/bert_squad_benchmark.py +1 -1

No files found.
--- a/official/benchmark/bert_squad_benchmark.py
+++ b/official/benchmark/bert_squad_benchmark.py
@@ -172,7 +172,7 @@ class BertSquadBenchmarkReal(BertSquadBenchmarkBase):
                                run_eagerly=False,
                                ds_type='mirrored'):
    """Runs the benchmark and reports various metrics."""
-    if FLAGS.train_batch_size <= 4:
+    if FLAGS.train_batch_size <= 4 or run_eagerly:
      FLAGS.input_meta_data_path = SQUAD_MEDIUM_INPUT_META_DATA_PATH
    else:
      FLAGS.input_meta_data_path = SQUAD_LONG_INPUT_META_DATA_PATH