Internal change

PiperOrigin-RevId: 271422100

Internal change
PiperOrigin-RevId: 271422100
d735ecf8 · David Chen · A. Unique TensorFlower · 127b158a · d735ecf8
Commit d735ecf8 authored Sep 26, 2019 by David Chen Committed by A. Unique TensorFlower Sep 26, 2019
Show whitespace changes
Inline Side-by-side

Showing with 2 additions and 2 deletions

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

No files found.
--- a/official/benchmark/bert_squad_benchmark.py
+++ b/official/benchmark/bert_squad_benchmark.py
@@ -82,7 +82,7 @@ class BertSquadBenchmarkBase(benchmark_utils.BertBenchmarkBase):
  @flagsaver.flagsaver
  def _train_squad(self, use_ds=True, run_eagerly=False):
    """Runs BERT SQuAD training."""
-    tf.enable_v2_behavior()
+    assert tf.version.VERSION.startswith('2.')
    input_meta_data = self._read_input_meta_data_from_file()
    strategy = self._get_distribution_strategy(use_ds)

@@ -95,7 +95,7 @@ class BertSquadBenchmarkBase(benchmark_utils.BertBenchmarkBase):
  @flagsaver.flagsaver
  def _evaluate_squad(self, use_ds=True):
    """Runs BERT SQuAD evaluation."""
-    tf.enable_v2_behavior()
+    assert tf.version.VERSION.startswith('2.')
    input_meta_data = self._read_input_meta_data_from_file()
    strategy = self._get_distribution_strategy(use_ds)