diff --git a/privacy/optimizers/gaussian_query_test.py b/privacy/optimizers/gaussian_query_test.py
index 7b5be7f..b66a442 100644
--- a/privacy/optimizers/gaussian_query_test.py
+++ b/privacy/optimizers/gaussian_query_test.py
@@ -20,6 +20,7 @@ from __future__ import print_function
 
 from absl.testing import parameterized
 import numpy as np
+from six.moves import xrange
 import tensorflow as tf
 
 from privacy.optimizers import gaussian_query
diff --git a/research/pate_2018/ICLR2018/rdp_bucketized.py b/research/pate_2018/ICLR2018/rdp_bucketized.py
index a5e7785..8527b46 100644
--- a/research/pate_2018/ICLR2018/rdp_bucketized.py
+++ b/research/pate_2018/ICLR2018/rdp_bucketized.py
@@ -40,11 +40,12 @@ sys.path.append('..')  # Main modules reside in the parent directory.
 
 from absl import app
 from absl import flags
+import core as pate
 import matplotlib
 matplotlib.use('TkAgg')
 import matplotlib.pyplot as plt  # pylint: disable=g-import-not-at-top
 import numpy as np
-import core as pate
+from six.moves import xrange
 
 plt.style.use('ggplot')