Skip to content

GitLab

Commit e64d4f56 authored by Hongkun Yu's avatar Hongkun Yu Committed by A. Unique TensorFlower
Browse files

Internal change 

PiperOrigin-RevId: 312765926
parent ff138931
Hide whitespace changes
Inline Side-by-side
Showing with 3 additions and 2 deletions
official/benchmark/nhnet_benchmark.py
View file @ e64d4f56
...@@ -39,14 +39,15 @@ FLAGS = flags.FLAGS ...@@ -39,14 +39,15 @@ FLAGS = flags.FLAGS
class NHNetBenchmark(perfzero_benchmark.PerfZeroBenchmark): class NHNetBenchmark(perfzero_benchmark.PerfZeroBenchmark):
"""Base benchmark class for NHNet.""" """Base benchmark class for NHNet."""
def __init__(self, output_dir=None, default_flags=None, tpu=None): def __init__(self, output_dir=None, default_flags=None, tpu=None, **kwargs):
self.default_flags = default_flags or {} self.default_flags = default_flags or {}
flag_methods = trainer.define_flags() flag_methods = trainer.define_flags()
super(NHNetBenchmark, self).__init__( super(NHNetBenchmark, self).__init__(
output_dir=output_dir, output_dir=output_dir,
default_flags=default_flags, default_flags=default_flags,
flag_methods=flag_methods, flag_methods=flag_methods,
tpu=tpu) tpu=tpu,
**kwargs)
def _report_benchmark(self, def _report_benchmark(self,
stats, stats,
......
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