Internal change

PiperOrigin-RevId: 297445627
tensorflow · Feb 26, 2020 · a208dd1 · a208dd1
1 parent 83f5681
commit a208dd1
Showing 1 changed file with 8 additions and 0 deletions.
diff --git a/official/benchmark/tfhub_memory_usage_benchmark.py b/official/benchmark/tfhub_memory_usage_benchmark.py
@@ -30,6 +30,14 @@
 class TfHubMemoryUsageBenchmark(PerfZeroBenchmark):
   """A benchmark measuring memory usage for a given TF Hub SavedModel."""
 
+  def __init__(self,
+               output_dir=None,
+               default_flags=None,
+               root_data_dir=None,
+               **kwargs):
+    super(TfHubMemoryUsageBenchmark, self).__init__(
+        output_dir=output_dir, default_flags=default_flags, **kwargs)
+
   def benchmark_memory_usage(self):
     start_time_sec = time.time()
     self.load_model()