diff --git a/research/pate_2018/ICLR2018/rdp_bucketized.py b/research/pate_2018/ICLR2018/rdp_bucketized.py
index a5e7785..f84ace1 100644
--- a/research/pate_2018/ICLR2018/rdp_bucketized.py
+++ b/research/pate_2018/ICLR2018/rdp_bucketized.py
@@ -36,6 +36,8 @@ import math
 import os
 import sys
 
+from six.moves import xrange
+
 sys.path.append('..')  # Main modules reside in the parent directory.
 
 from absl import app