diff --git a/create_pretraining_data.py b/create_pretraining_data.py
index 5340d96ae..1daab7d9a 100644
--- a/create_pretraining_data.py
+++ b/create_pretraining_data.py
@@ -22,6 +22,9 @@
 import random
 import tokenization
 import tensorflow as tf
+if(tf.__version__.startswith("2")):
+    tf.gfile = tf.io.gfile
+    tf.flags = tf.compat.v1.flags
 
 flags = tf.flags