diff --git a/classifier_utils.py b/classifier_utils.py
index 0dbfaadb..96cfef1a 100644
--- a/classifier_utils.py
+++ b/classifier_utils.py
@@ -24,7 +24,7 @@
 import modeling
 import optimization
 import tokenization
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 from tensorflow.contrib import data as contrib_data
 from tensorflow.contrib import metrics as contrib_metrics
 from tensorflow.contrib import tpu as contrib_tpu
diff --git a/lamb_optimizer.py b/lamb_optimizer.py
index 074bfb71..a7222d2c 100644
--- a/lamb_optimizer.py
+++ b/lamb_optimizer.py
@@ -21,7 +21,7 @@
 
 import re
 import six
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 # pylint: disable=g-direct-tensorflow-import
 from tensorflow.python.ops import array_ops
diff --git a/modeling.py b/modeling.py
index 93997d42..c9a281a4 100644
--- a/modeling.py
+++ b/modeling.py
@@ -30,7 +30,7 @@
 import numpy as np
 import six
 from six.moves import range
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 from tensorflow.contrib import layers as contrib_layers
 
 
diff --git a/modeling_test.py b/modeling_test.py
index 8815dfb5..567c8023 100644
--- a/modeling_test.py
+++ b/modeling_test.py
@@ -25,7 +25,7 @@
 import modeling
 import six
 from six.moves import range
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 class AlbertModelTest(tf.test.TestCase):
diff --git a/optimization.py b/optimization.py
index e34e7245..0df92cc1 100644
--- a/optimization.py
+++ b/optimization.py
@@ -22,7 +22,7 @@
 import lamb_optimizer
 import six
 from six.moves import zip
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 from tensorflow.contrib import tpu as contrib_tpu
 
 
diff --git a/optimization_test.py b/optimization_test.py
index c035d219..156b5a05 100644
--- a/optimization_test.py
+++ b/optimization_test.py
@@ -19,7 +19,7 @@
 import optimization
 from six.moves import range
 from six.moves import zip
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 class OptimizationTest(tf.test.TestCase):
diff --git a/squad_utils.py b/squad_utils.py
index e357fb1c..6f4ace84 100644
--- a/squad_utils.py
+++ b/squad_utils.py
@@ -32,7 +32,7 @@
 import six
 from six.moves import map
 from six.moves import range
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 from tensorflow.contrib import data as contrib_data
 from tensorflow.contrib import layers as contrib_layers
 from tensorflow.contrib import tpu as contrib_tpu
diff --git a/tokenization.py b/tokenization.py
index a04a11c4..4cad9527 100644
--- a/tokenization.py
+++ b/tokenization.py
@@ -25,7 +25,7 @@
 import unicodedata
 import six
 from six.moves import range
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 import sentencepiece as spm
 
 SPIECE_UNDERLINE = u"▁".encode("utf-8")
diff --git a/tokenization_test.py b/tokenization_test.py
index 3aa1f34e..99313a0c 100644
--- a/tokenization_test.py
+++ b/tokenization_test.py
@@ -20,7 +20,7 @@
 import tempfile
 import tokenization
 import six
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 class TokenizationTest(tf.test.TestCase):