Internal change

PiperOrigin-RevId: 308880055

official/benchmark/bert_benchmark_utils.py

@@ -67,8 +67,9 @@ class BertBenchmarkBase(PerfZeroBenchmark):
   """Base class to hold methods common to test classes."""
   local_flags = None
 
-  def __init__(self, output_dir=None, tpu=None):
-    super(BertBenchmarkBase, self).__init__(output_dir=output_dir, tpu=tpu)
+  def __init__(self, output_dir=None, tpu=None, **kwargs):
+    super(BertBenchmarkBase, self).__init__(
+        output_dir=output_dir, tpu=tpu, **kwargs)
     self.num_gpus = 8
     self.timer_callback = None
 

official/benchmark/bert_pretrain_benchmark.py

@@ -51,15 +51,17 @@ class BertPretrainAccuracyBenchmark(bert_benchmark_utils.BertBenchmarkBase):
 
   def __init__(self,
                output_dir: Optional[str] = None,
-               tpu: Optional[str] = None):
+               tpu: Optional[str] = None,
+               **kwargs):
     """Inits BertPretrainAccuracyBenchmark class.
 
     Args:
       output_dir: Directory where to output e.g. log files
       tpu: TPU name to use in a TPU benchmark.
+      **kwargs: Additional keyword arguments.
     """
     super(BertPretrainAccuracyBenchmark, self).__init__(
-        output_dir=output_dir, tpu=tpu)
+        output_dir=output_dir, tpu=tpu, **kwargs)
 
   @benchmark_wrappers.enable_runtime_flags
   def _run_and_report_benchmark(self, summary_path: str):