diff --git a/tensorflow_privacy/privacy/optimizers/dp_optimizer_eager_test.py b/tensorflow_privacy/privacy/optimizers/dp_optimizer_eager_test.py
index e993842..489aebe 100644
--- a/tensorflow_privacy/privacy/optimizers/dp_optimizer_eager_test.py
+++ b/tensorflow_privacy/privacy/optimizers/dp_optimizer_eager_test.py
@@ -19,6 +19,7 @@ from __future__ import print_function
 
 from absl.testing import parameterized
 import numpy as np
+from six.moves import range
 import tensorflow.compat.v1 as tf
 
 from tensorflow_privacy.privacy.analysis import privacy_ledger