diff --git a/tensorflow_privacy/privacy/optimizers/dp_optimizer_vectorized_test.py b/tensorflow_privacy/privacy/optimizers/dp_optimizer_vectorized_test.py
index 83aaf48..d254817 100644
--- a/tensorflow_privacy/privacy/optimizers/dp_optimizer_vectorized_test.py
+++ b/tensorflow_privacy/privacy/optimizers/dp_optimizer_vectorized_test.py
@@ -20,6 +20,7 @@ from __future__ import print_function
 from absl.testing import parameterized
 import mock
 import numpy as np
+from six.moves import range
 import tensorflow.compat.v1 as tf
 
 from tensorflow_privacy.privacy.optimizers import dp_optimizer_vectorized