diff --git a/setup.py b/setup.py
index 0c6df6f73..b68e4ce6a 100644
--- a/setup.py
+++ b/setup.py
@@ -5,7 +5,7 @@
 
 setup(
     name='tensor2tensor',
-    version='1.15.1',
+    version='1.15.2',
     description='Tensor2Tensor',
     long_description=(
         'Tensor2Tensor, or T2T for short, is a library of '
@@ -66,7 +66,6 @@
         'tensorflow-datasets',
         'tensorflow-gan',
         'tensorflow-probability==0.7.0',
-        'tf_slim',
         'tqdm',
     ],
     extras_require={
diff --git a/tensor2tensor/data_generators/allen_brain.py b/tensor2tensor/data_generators/allen_brain.py
index 2a738885a..2f58de3e9 100644
--- a/tensor2tensor/data_generators/allen_brain.py
+++ b/tensor2tensor/data_generators/allen_brain.py
@@ -44,7 +44,6 @@
 from tensor2tensor.utils import registry
 
 import tensorflow as tf
-import tf_slim as slim
 
 _BASE_EXAMPLE_IMAGE_SIZE = 64
 
@@ -351,7 +350,7 @@ def example_reading_spec(self):
 
     data_items_to_decoders = {
         "targets":
-            slim.tfexample_decoder.Image(
+            tf.contrib.slim.tfexample_decoder.Image(
                 image_key="image/encoded",
                 format_key="image/format",
                 channels=self.num_channels),
diff --git a/tensor2tensor/data_generators/bair_robot_pushing.py b/tensor2tensor/data_generators/bair_robot_pushing.py
index b4e094dfb..eadd3275f 100644
--- a/tensor2tensor/data_generators/bair_robot_pushing.py
+++ b/tensor2tensor/data_generators/bair_robot_pushing.py
@@ -36,7 +36,6 @@
 from tensor2tensor.utils import registry
 
 import tensorflow as tf
-import tf_slim as slim
 
 DATA_URL = (
     "http://rail.eecs.berkeley.edu/datasets/bair_robot_pushing_dataset_v0.tar")
@@ -103,7 +102,7 @@ def extra_reading_spec(self):
         "frame_number": tf.FixedLenFeature([1], tf.int64),
     }
     decoders = {
-        "frame_number": slim.tfexample_decoder.Tensor(
+        "frame_number": tf.contrib.slim.tfexample_decoder.Tensor(
             tensor_key="frame_number"),
     }
     return data_fields, decoders
@@ -188,9 +187,9 @@ def extra_reading_spec(self):
         "action": tf.FixedLenFeature([4], tf.float32),
     }
     decoders = {
-        "frame_number": slim.tfexample_decoder.Tensor(
+        "frame_number": tf.contrib.slim.tfexample_decoder.Tensor(
             tensor_key="frame_number"),
-        "action": slim.tfexample_decoder.Tensor(tensor_key="action"),
+        "action": tf.contrib.slim.tfexample_decoder.Tensor(tensor_key="action"),
     }
     return data_fields, decoders
 
diff --git a/tensor2tensor/data_generators/fsns.py b/tensor2tensor/data_generators/fsns.py
index 4792263d7..8dff69547 100644
--- a/tensor2tensor/data_generators/fsns.py
+++ b/tensor2tensor/data_generators/fsns.py
@@ -28,7 +28,6 @@
 from tensor2tensor.utils import registry
 
 import tensorflow as tf
-import tf_slim as slim
 
 
 @registry.register_problem
@@ -77,5 +76,5 @@ def example_reading_spec(self):
         super(ImageFSNS, self).example_reading_spec())
     data_fields[label_key] = tf.VarLenFeature(tf.int64)
     data_items_to_decoders[
-        "targets"] = slim.tfexample_decoder.Tensor(label_key)
+        "targets"] = tf.contrib.slim.tfexample_decoder.Tensor(label_key)
     return data_fields, data_items_to_decoders
diff --git a/tensor2tensor/data_generators/gym_env.py b/tensor2tensor/data_generators/gym_env.py
index c325772cf..67aedef9b 100644
--- a/tensor2tensor/data_generators/gym_env.py
+++ b/tensor2tensor/data_generators/gym_env.py
@@ -36,7 +36,6 @@
 from tensor2tensor.utils import registry
 
 import tensorflow as tf
-import tf_slim as slim
 
 
 Frame = collections.namedtuple(
@@ -378,7 +377,7 @@ def extra_reading_spec(self):
         name: tf.FixedLenFeature([1], tf.int64) for name in field_names
     }
     decoders = {
-        name: slim.tfexample_decoder.Tensor(tensor_key=name)
+        name: tf.contrib.slim.tfexample_decoder.Tensor(tensor_key=name)
         for name in field_names
     }
     return (data_fields, decoders)
diff --git a/tensor2tensor/data_generators/image_utils.py b/tensor2tensor/data_generators/image_utils.py
index 23c38c543..61766a89e 100644
--- a/tensor2tensor/data_generators/image_utils.py
+++ b/tensor2tensor/data_generators/image_utils.py
@@ -31,7 +31,6 @@
 from tensor2tensor.utils import metrics
 
 import tensorflow as tf
-import tf_slim as slim
 
 
 def matplotlib_pyplot():
@@ -173,7 +172,7 @@ def example_reading_spec(self):
 
     data_items_to_decoders = {
         "inputs":
-            slim.tfexample_decoder.Image(
+            tf.contrib.slim.tfexample_decoder.Image(
                 image_key="image/encoded",
                 format_key="image/format",
                 channels=self.num_channels),
@@ -240,7 +239,7 @@ def example_reading_spec(self):
     data_fields[label_key] = tf.FixedLenFeature((1,), tf.int64)
 
     data_items_to_decoders[
-        "targets"] = slim.tfexample_decoder.Tensor(label_key)
+        "targets"] = tf.contrib.slim.tfexample_decoder.Tensor(label_key)
     return data_fields, data_items_to_decoders
 
   def hparams(self, defaults, unused_model_hparams):
@@ -344,7 +343,7 @@ def example_reading_spec(self):
         super(Image2TextProblem, self).example_reading_spec())
     data_fields[label_key] = tf.VarLenFeature(tf.int64)
     data_items_to_decoders[
-        "targets"] = slim.tfexample_decoder.Tensor(label_key)
+        "targets"] = tf.contrib.slim.tfexample_decoder.Tensor(label_key)
     return data_fields, data_items_to_decoders
 
   def feature_encoders(self, data_dir):
diff --git a/tensor2tensor/data_generators/moving_mnist.py b/tensor2tensor/data_generators/moving_mnist.py
index d8841c741..507207623 100644
--- a/tensor2tensor/data_generators/moving_mnist.py
+++ b/tensor2tensor/data_generators/moving_mnist.py
@@ -37,7 +37,6 @@
 import tensorflow as tf
 import tensorflow_datasets as tfds
 from tensorflow_datasets.video import moving_sequence
-import tf_slim as slim
 
 
 DATA_URL = (
@@ -95,7 +94,7 @@ def extra_reading_spec(self):
         "frame_number": tf.FixedLenFeature([1], tf.int64),
     }
     decoders = {
-        "frame_number": slim.tfexample_decoder.Tensor(
+        "frame_number": tf.contrib.slim.tfexample_decoder.Tensor(
             tensor_key="frame_number"),
     }
     return data_fields, decoders
diff --git a/tensor2tensor/data_generators/problem.py b/tensor2tensor/data_generators/problem.py
index c1ff4ef4f..aca91f31f 100644
--- a/tensor2tensor/data_generators/problem.py
+++ b/tensor2tensor/data_generators/problem.py
@@ -32,8 +32,13 @@
 from tensor2tensor.utils import mlperf_log
 
 import tensorflow as tf
-import tf_slim as slim
-from tensorflow.contrib.tpu.python.tpu import tpu_config
+# pylint: disable=g-import-not-at-top
+try:
+  from tensorflow.contrib.tpu.python.tpu import tpu_config
+except ImportError:
+  # TF 2.0 doesn't ship with contrib.
+  tpu_config = None
+# pylint: enable=g-import-not-at-top
 
 
 
@@ -199,7 +204,7 @@ class Problem(object):
         - Mutate defaults as needed
     * example_reading_spec
         - Specify the names and types of the features on disk.
-        - Specify slim.tfexample_decoder
+        - Specify tf.contrib.slim.tfexample_decoder
     * preprocess_example(example, mode, hparams)
         - Preprocess the example feature dict from feature name to Tensor or
           SparseTensor.
@@ -643,7 +648,7 @@ def dataset(self,
 
     data_filepattern = self.filepattern(data_dir, dataset_split, shard=shard)
     tf.logging.info("Reading data files from %s", data_filepattern)
-    data_files = sorted(slim.parallel_reader.get_data_files(
+    data_files = sorted(tf.contrib.slim.parallel_reader.get_data_files(
         data_filepattern))
 
     # Functions used in dataset transforms below. `filenames` can be either a
@@ -706,12 +711,12 @@ def decode_example(self, serialized_example):
     data_fields["batch_prediction_key"] = tf.FixedLenFeature([1], tf.int64, 0)
     if data_items_to_decoders is None:
       data_items_to_decoders = {
-          field: slim.tfexample_decoder.Tensor(field)
+          field: tf.contrib.slim.tfexample_decoder.Tensor(field)
           for field in data_fields
       }
 
-    decoder = slim.tfexample_decoder.TFExampleDecoder(data_fields,
-                                                      data_items_to_decoders)
+    decoder = tf.contrib.slim.tfexample_decoder.TFExampleDecoder(
+        data_fields, data_items_to_decoders)
 
     decode_items = list(sorted(data_items_to_decoders))
     decoded = decoder.decode(serialized_example, items=decode_items)
diff --git a/tensor2tensor/data_generators/style_transfer.py b/tensor2tensor/data_generators/style_transfer.py
index cdd9ceacc..146793666 100644
--- a/tensor2tensor/data_generators/style_transfer.py
+++ b/tensor2tensor/data_generators/style_transfer.py
@@ -32,16 +32,12 @@
 from tensor2tensor.data_generators import text_problems
 from tensor2tensor.utils import registry
 
-import tensorflow as tf
 
-logger = tf.logging
+# Modern-Shakespeare corpus is consisted of:
+# - 18,395 parallel sentences for training (train set),
+# - 1,218 parallel sentences for evaluation (dev set),
+# - 1,462 parallel sentence for testing (test set).
 
-"""
-Modern-Shakespeare corpus is consisted of:
-- 18,395 parallel sentences for training (train set),
-- 1,218 parallel sentences for evaluation (dev set),
-- 1,462 parallel sentence for testing (test set).
-"""
 
 _SHAKESPEARE_MODERN_TRAIN_DATASET = [[
     "https://github.com/tlatkowski/st/raw/master/shakespeare.train.tgz",
diff --git a/tensor2tensor/data_generators/translate.py b/tensor2tensor/data_generators/translate.py
index 8ac2f2575..05e07d7c3 100644
--- a/tensor2tensor/data_generators/translate.py
+++ b/tensor2tensor/data_generators/translate.py
@@ -32,7 +32,6 @@
 from tensor2tensor.utils import mlperf_log
 
 import tensorflow as tf
-import tf_slim as slim
 
 
 class TranslateProblem(text_problems.Text2TextProblem):
@@ -277,8 +276,8 @@ def example_reading_spec(self):
 
     # hack: ignoring true targets and putting dist_targets in targets
     data_items_to_decoders = {
-        "inputs": slim.tfexample_decoder.Tensor("inputs"),
-        "targets": slim.tfexample_decoder.Tensor("dist_targets"),
+        "inputs": tf.contrib.slim.tfexample_decoder.Tensor("inputs"),
+        "targets": tf.contrib.slim.tfexample_decoder.Tensor("dist_targets"),
     }
 
     return (data_fields, data_items_to_decoders)
diff --git a/tensor2tensor/data_generators/video_generated.py b/tensor2tensor/data_generators/video_generated.py
index e4f2e5661..aa0e5bcf3 100644
--- a/tensor2tensor/data_generators/video_generated.py
+++ b/tensor2tensor/data_generators/video_generated.py
@@ -28,7 +28,6 @@
 from tensor2tensor.utils import registry
 
 import tensorflow as tf
-import tf_slim as slim
 
 try:
   import matplotlib  # pylint: disable=g-import-not-at-top
@@ -86,7 +85,7 @@ def extra_reading_spec(self):
         "frame_number": tf.FixedLenFeature([1], tf.int64),
     }
     decoders = {
-        "frame_number": slim.tfexample_decoder.Tensor(
+        "frame_number": tf.contrib.slim.tfexample_decoder.Tensor(
             tensor_key="frame_number"),
     }
     return data_fields, decoders
diff --git a/tensor2tensor/data_generators/video_utils.py b/tensor2tensor/data_generators/video_utils.py
index 4d3cec63c..ae2d2b313 100644
--- a/tensor2tensor/data_generators/video_utils.py
+++ b/tensor2tensor/data_generators/video_utils.py
@@ -35,7 +35,6 @@
 from tensor2tensor.utils import metrics
 from tensor2tensor.utils import video_metrics
 import tensorflow as tf
-import tf_slim as slim
 
 
 FLAGS = flags.FLAGS
@@ -385,7 +384,7 @@ def example_reading_spec(self):
 
     data_items_to_decoders = {
         "frame":
-            slim.tfexample_decoder.Image(
+            tf.contrib.slim.tfexample_decoder.Image(
                 image_key="image/encoded",
                 format_key="image/format",
                 shape=[self.frame_height, self.frame_width, self.num_channels],
@@ -677,7 +676,7 @@ def example_reading_spec(self):
 
     data_items_to_decoders = {
         "inputs":
-            slim.tfexample_decoder.Image(
+            tf.contrib.slim.tfexample_decoder.Image(
                 image_key="image/encoded",
                 format_key="image/format",
                 channels=self.num_channels),
@@ -767,7 +766,7 @@ def example_reading_spec(self):
         super(Video2ClassProblem, self).example_reading_spec())
     data_fields[label_key] = tf.FixedLenFeature((1,), tf.int64)
     data_items_to_decoders[
-        "targets"] = slim.tfexample_decoder.Tensor(label_key)
+        "targets"] = tf.contrib.slim.tfexample_decoder.Tensor(label_key)
     return data_fields, data_items_to_decoders
 
   def hparams(self, defaults, unused_model_hparams):
diff --git a/tensor2tensor/data_generators/vqa.py b/tensor2tensor/data_generators/vqa.py
index e0c75d388..8a39f85f6 100644
--- a/tensor2tensor/data_generators/vqa.py
+++ b/tensor2tensor/data_generators/vqa.py
@@ -39,7 +39,6 @@
 from tensor2tensor.utils import registry
 
 import tensorflow as tf
-import tf_slim as slim
 
 
 def _get_vqa_v2_annotations(directory,
@@ -218,10 +217,10 @@ def example_reading_spec(self):
         (), tf.int64, allow_missing=True)
 
     data_items_to_decoders[
-        "question"] = slim.tfexample_decoder.Tensor(
+        "question"] = tf.contrib.slim.tfexample_decoder.Tensor(
             "image/question")
     data_items_to_decoders[
-        "targets"] = slim.tfexample_decoder.Tensor(
+        "targets"] = tf.contrib.slim.tfexample_decoder.Tensor(
             "image/answer")
     return data_fields, data_items_to_decoders
 
@@ -339,23 +338,23 @@ def example_reading_spec(self):
         (), tf.int64, allow_missing=True)
 
     data_items_to_decoders[
-        "inputs"] = slim.tfexample_decoder.Tensor(
+        "inputs"] = tf.contrib.slim.tfexample_decoder.Tensor(
             "image/feature")
     data_items_to_decoders[
-        "question_id"] = slim.tfexample_decoder.Tensor(
+        "question_id"] = tf.contrib.slim.tfexample_decoder.Tensor(
             "image/question_id")
     data_items_to_decoders[
-        "image_id"] = slim.tfexample_decoder.Tensor(
+        "image_id"] = tf.contrib.slim.tfexample_decoder.Tensor(
             "image/image_id")
 
     data_items_to_decoders[
-        "spatial_feature"] = slim.tfexample_decoder.Tensor(
+        "spatial_feature"] = tf.contrib.slim.tfexample_decoder.Tensor(
             "image/spatial_feature")
     data_items_to_decoders[
-        "question"] = slim.tfexample_decoder.Tensor(
+        "question"] = tf.contrib.slim.tfexample_decoder.Tensor(
             "image/question")
     data_items_to_decoders[
-        "targets"] = slim.tfexample_decoder.Tensor(
+        "targets"] = tf.contrib.slim.tfexample_decoder.Tensor(
             "image/answer")
 
     return data_fields, data_items_to_decoders
diff --git a/tensor2tensor/envs/env_problem.py b/tensor2tensor/envs/env_problem.py
index 7c02f43fd..43afb1a75 100644
--- a/tensor2tensor/envs/env_problem.py
+++ b/tensor2tensor/envs/env_problem.py
@@ -33,7 +33,6 @@
 from tensor2tensor.envs import trajectory
 from tensor2tensor.layers import modalities
 import tensorflow as tf
-import tf_slim as slim
 
 # Names for data fields in stored tf.Examples.
 TIMESTEP_FIELD = "timestep"
@@ -477,7 +476,7 @@ def example_reading_spec(self):
     }
 
     data_items_to_decoders = {
-        field: slim.tfexample_decoder.Tensor(field)
+        field: tf.contrib.slim.tfexample_decoder.Tensor(field)
         for field in data_fields
     }
 
diff --git a/tensor2tensor/envs/rendered_env_problem.py b/tensor2tensor/envs/rendered_env_problem.py
index 82eb2fa71..92a7a7e54 100644
--- a/tensor2tensor/envs/rendered_env_problem.py
+++ b/tensor2tensor/envs/rendered_env_problem.py
@@ -25,7 +25,6 @@
 from tensor2tensor.envs import env_problem
 from tensor2tensor.envs import gym_env_problem
 import tensorflow as tf
-import tf_slim as slim
 
 _IMAGE_ENCODED_FIELD = "image/encoded"
 _IMAGE_FORMAT_FIELD = "image/format"
@@ -81,7 +80,7 @@ def example_reading_spec(self):
     # Add frame number spec and decoder.
     env_fields[_FRAME_NUMBER_FIELD] = tf.FixedLenFeature((1,), tf.int64)
     env_decoders[
-        _FRAME_NUMBER_FIELD] = slim.tfexample_decoder.Tensor(
+        _FRAME_NUMBER_FIELD] = tf.contrib.slim.tfexample_decoder.Tensor(
             _FRAME_NUMBER_FIELD)
 
     # Add video fields and decoders
diff --git a/tensor2tensor/envs/trajectory.py b/tensor2tensor/envs/trajectory.py
index cf4705a1c..e980db4d5 100644
--- a/tensor2tensor/envs/trajectory.py
+++ b/tensor2tensor/envs/trajectory.py
@@ -31,7 +31,7 @@
 import cloudpickle
 import numpy as np
 from tensor2tensor.envs import time_step
-from tensorflow.io import gfile
+import tensorflow as tf
 
 TRAJECTORY_FILE_FORMAT = r"trajectory_epoch_{epoch}_env_id_{env_id}_temperature_{temperature}_r_{r}.pkl"
 
@@ -528,7 +528,7 @@ def load_from_directory(trajectory_dir,
         r="*",
     )
 
-    trajectory_files = gfile.glob(
+    trajectory_files = tf.io.gfile.glob(
         os.path.join(trajectory_dir, trajectory_file_glob))
 
     if n_trajectories:
@@ -543,7 +543,7 @@ def load_from_directory(trajectory_dir,
         time.sleep(sleep_time_secs)
         max_tries -= 1
         sleep_time_secs = min(10.0, sleep_time_secs * 2)
-        trajectory_files = gfile.glob(
+        trajectory_files = tf.io.gfile.glob(
             os.path.join(trajectory_dir, trajectory_file_glob))
 
       # We can't get the required number of files and we can't up-sample either.
@@ -557,7 +557,7 @@ def load_from_directory(trajectory_dir,
     # We read and load all the files, revisit if this becomes a problem.
     trajectories_buffer = []
     for trajectory_file in trajectory_files:
-      with gfile.GFile(trajectory_file, "rb") as f:
+      with tf.io.gfile.GFile(trajectory_file, "rb") as f:
         trajectory = get_pickle_module().load(f)
         assert isinstance(trajectory, Trajectory)
         trajectories_buffer.append(trajectory)
diff --git a/tensor2tensor/layers/common_attention.py b/tensor2tensor/layers/common_attention.py
index f3c43443e..a1589ed5d 100644
--- a/tensor2tensor/layers/common_attention.py
+++ b/tensor2tensor/layers/common_attention.py
@@ -951,7 +951,8 @@ def attention_bias_ignore_padding(memory_padding):
 
 
 @expert_utils.add_name_scope()
-def attention_bias_to_padding(attention_bias, cast_fn=tf.to_float):
+def attention_bias_to_padding(attention_bias,
+                              cast_fn=(lambda x: tf.cast(x, tf.float32))):
   """Inverse of attention_bias_ignore_padding().
 
   Args:
diff --git a/tensor2tensor/layers/common_layers.py b/tensor2tensor/layers/common_layers.py
index 8eb6a8d19..aae4a6b4c 100644
--- a/tensor2tensor/layers/common_layers.py
+++ b/tensor2tensor/layers/common_layers.py
@@ -23,6 +23,7 @@
 import functools
 import math
 
+from absl import logging
 import numpy as np
 from six.moves import range  # pylint: disable=redefined-builtin
 
@@ -44,11 +45,15 @@ def layers():
   global _cached_layers
   if _cached_layers is not None:
     return _cached_layers
-  layers_module = tf.layers
+  layers_module = None
+  try:
+    layers_module = tf.layers
+  except AttributeError:
+    logging.info("Cannot access tf.layers, trying TF2 layers.")
   try:
     from tensorflow.python import tf2  # pylint: disable=g-direct-tensorflow-import,g-import-not-at-top
     if tf2.enabled():
-      tf.logging.info("Running in V2 mode, using Keras layers.")
+      logging.info("Running in V2 mode, using Keras layers.")
       layers_module = tf.keras.layers
   except ImportError:
     pass
diff --git a/tensor2tensor/layers/common_video.py b/tensor2tensor/layers/common_video.py
index 0d2a1af4c..d564db7fe 100644
--- a/tensor2tensor/layers/common_video.py
+++ b/tensor2tensor/layers/common_video.py
@@ -33,9 +33,12 @@
 except ImportError:
   distribute_summary_op_util = summary_op_util
 
-
-tfl = tf.layers
-tfcl = tf.contrib.layers
+tfl = common_layers.layers()
+tfcl = None
+try:
+  tfcl = tf.contrib.layers
+except AttributeError:
+  pass
 
 
 def swap_time_and_batch_axes(inputs):
diff --git a/tensor2tensor/utils/data_reader.py b/tensor2tensor/utils/data_reader.py
index e80aa6dc4..c65878163 100644
--- a/tensor2tensor/utils/data_reader.py
+++ b/tensor2tensor/utils/data_reader.py
@@ -28,7 +28,6 @@
 from tensor2tensor.utils import mlperf_log
 
 import tensorflow as tf
-import tf_slim as slim
 
 
 def cast_ints_to_int32(features):
@@ -380,7 +379,7 @@ def define_shapes(example):
     dataset = dataset.repeat()
 
   if is_training and skip_random_fraction_when_training:
-    data_files = slim.parallel_reader.get_data_files(filepattern)
+    data_files = tf.contrib.slim.parallel_reader.get_data_files(filepattern)
     #  In continuous_train_and_eval when switching between train and
     #  eval, this input_fn method gets called multiple times and it
     #  would give you the exact same samples from the last call
diff --git a/tensor2tensor/utils/metrics.py b/tensor2tensor/utils/metrics.py
index fcdf1aab2..59a54133b 100644
--- a/tensor2tensor/utils/metrics.py
+++ b/tensor2tensor/utils/metrics.py
@@ -29,7 +29,6 @@
 
 import tensorflow as tf
 
-from tensorflow.contrib.eager.python import tfe
 from tensorflow.python.util import tf_inspect as inspect
 
 
@@ -786,6 +785,9 @@ def create_eager_metrics_internal(metric_fns,
     (accum_fn(predictions, targets) => None,
      result_fn() => dict<str metric_name, float avg_val>
   """
+
+  from tensorflow.contrib.eager.python import tfe  # pylint: disable=g-import-not-at-top
+
   tfe_metrics = {}
 
   for name in metric_fns:
diff --git a/tensor2tensor/utils/trainer_lib.py b/tensor2tensor/utils/trainer_lib.py
index 5a108f53c..36ffe363b 100644
--- a/tensor2tensor/utils/trainer_lib.py
+++ b/tensor2tensor/utils/trainer_lib.py
@@ -36,7 +36,6 @@
 
 import tensorflow as tf
 
-from tensorflow.contrib.tpu.python.tpu import tpu_estimator
 from tensorflow.core.protobuf import rewriter_config_pb2
 from tensorflow.python import debug
 
@@ -297,6 +296,7 @@ def create_estimator(model_name,
 
   del use_xla
   if use_tpu or use_tpu_estimator:
+    from tensorflow.contrib.tpu.python.tpu import tpu_estimator  # pylint: disable=g-import-not-at-top
     problem = hparams.problem
     batch_size = (
         problem.tpu_batch_size_per_shard(hparams) *