From a208dd14ea02ad672517e7568f6cd2e1a6525a01 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Wed, 26 Feb 2020 14:03:54 -0800
Subject: [PATCH] Internal change

PiperOrigin-RevId: 297445627
---
 official/benchmark/tfhub_memory_usage_benchmark.py | 8 ++++++++
 1 file changed, 8 insertions(+)

diff --git a/official/benchmark/tfhub_memory_usage_benchmark.py b/official/benchmark/tfhub_memory_usage_benchmark.py
index 9d4861c0177..2c0a2fceef9 100644
--- 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()