diff --git a/generate_tfrecord.py b/generate_tfrecord.py
index bb7379a5fb4..7936492cd11 100644
--- a/generate_tfrecord.py
+++ b/generate_tfrecord.py
@@ -15,7 +15,7 @@
 import os
 import io
 import pandas as pd
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 import sys
 
 sys.path.append("../../models/research")