Skip to content

GitLab

Commit 458ea726 authored by Zongwei Zhou's avatar Zongwei Zhou Committed by A. Unique TensorFlower
Browse files

Fix BERT SQUAD eager tests 

PiperOrigin-RevId: 292947055
parent f4147e22
Hide whitespace changes
Inline Side-by-side
Showing with 2 additions and 4 deletions
official/benchmark/bert_squad_benchmark.py
View file @ 458ea726
......@@ -190,11 +190,10 @@ class BertSquadBenchmarkReal(BertSquadBenchmarkBase):
self._setup()
self.num_gpus = 1
FLAGS.run_eagerly = True
FLAGS.model_dir = self._get_model_dir('benchmark_1_gpu_squad_eager')
FLAGS.train_batch_size = 2
self._run_and_report_benchmark()
self._run_and_report_benchmark(run_eagerly=True)
def benchmark_1_gpu_xla(self):
"""Tests BERT SQuAD model performance with 1 GPU with XLA."""
......@@ -265,13 +264,12 @@ class BertSquadBenchmarkReal(BertSquadBenchmarkBase):
self._setup()
self.num_gpus = 1
FLAGS.run_eagerly = True
FLAGS.model_dir = self._get_model_dir('benchmark_1_gpu_squad_fp16_eager')
FLAGS.train_batch_size = 4
FLAGS.dtype = 'fp16'
FLAGS.loss_scale = 'dynamic'
self._run_and_report_benchmark()
self._run_and_report_benchmark(run_eagerly=True)
def benchmark_1_gpu_fp16(self):
"""Tests BERT SQuAD model performance with 1 GPU and FP16."""
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment