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