diff --git a/classifier_utils.py b/classifier_utils.py
index 0edcbd7..b5579df 100644
--- a/classifier_utils.py
+++ b/classifier_utils.py
@@ -26,7 +26,6 @@
 from albert import tokenization
 import tensorflow.compat.v1 as tf
 from tensorflow.compat.v1 import estimator as tf_estimator
-from tensorflow.contrib import data as contrib_data
 from tensorflow.contrib import metrics as contrib_metrics
 from tensorflow.contrib import tpu as contrib_tpu
 
@@ -739,7 +738,7 @@ def input_fn(params):
       d = d.shuffle(buffer_size=100)
 
     d = d.apply(
-        contrib_data.map_and_batch(
+        tf.data.experimental.map_and_batch(
             lambda record: _decode_record(record, name_to_features),
             batch_size=batch_size,
             drop_remainder=drop_remainder))
diff --git a/squad_utils.py b/squad_utils.py
index 3e2a740..9696609 100644
--- a/squad_utils.py
+++ b/squad_utils.py
@@ -33,7 +33,6 @@
 from six.moves import range
 import tensorflow.compat.v1 as tf
 from tensorflow.compat.v1 import estimator as tf_estimator
-from tensorflow.contrib import data as contrib_data
 from tensorflow.contrib import layers as contrib_layers
 from tensorflow.contrib import tpu as contrib_tpu
 
@@ -696,7 +695,7 @@ def input_fn(params):
       d = d.shuffle(buffer_size=100)
 
     d = d.apply(
-        contrib_data.map_and_batch(
+        tf.data.experimental.map_and_batch(
             lambda record: _decode_record(record, name_to_features),
             batch_size=batch_size,
             drop_remainder=drop_remainder))