Initialize context handle before enable_run_metadata.

PiperOrigin-RevId: 181770918

1 files changed, 10 insertions, 0 deletions

diff --git a/tensorflow/python/profiler/model_analyzer_test.py b/tensorflow/python/profiler/model_analyzer_test.py
index 14ad9e51a1..0ad4e0ef17 100644
--- a/tensorflow/python/profiler/model_analyzer_test.py
+++ b/tensorflow/python/profiler/model_analyzer_test.py
@@ -29,6 +29,7 @@ from tensorflow.core.profiler import profile_pb2
 from tensorflow.core.protobuf import config_pb2
 from tensorflow.core.protobuf import rewriter_config_pb2
 from tensorflow.python.client import session
+from tensorflow.python.eager import context
 from tensorflow.python.framework import dtypes
 from tensorflow.python.framework import ops
 from tensorflow.python.framework import test_util
@@ -762,6 +763,15 @@ class PrintModelAnalysisTest(test.TestCase):
           sess.graph, run_meta=run_metadata, cmd='scope', options=options)
       self.assertGreater(ret_pb.total_requested_bytes, 1000000)
 
+  def testEager(self):
+    ops.reset_default_graph()
+    with context.eager_mode():
+      context.enable_run_metadata()
+      lib.BuildSmallModel()
+      run_meta = context.export_run_metadata()
+      self.assertTrue('Conv2D' in '%s' % run_meta)
+      context.disable_run_metadata()
+
 
 if __name__ == '__main__':
   test.main()