Use tf.compat.v1 in cloud_tpu/models/common.

will-cromar · allenwang28 · commit ae62fcd319b2 · 2019-10-08T08:21:46.000-07:00
PiperOrigin-RevId: 272715691
diff --git a/models/common/inference_warmup.py b/models/common/inference_warmup.py
@@ -22,7 +22,7 @@
 import io
 import os
 import numpy as np
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 def _encode_image(image_array, fmt):
diff --git a/models/common/tpu_profiler_hook.py b/models/common/tpu_profiler_hook.py
@@ -21,7 +21,7 @@
 
 import subprocess
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 class TPUProfilerHook(tf.train.SessionRunHook):
