diff --git a/python/array_record_data_source.py b/python/array_record_data_source.py
index 4a6aed3..9ffe21a 100644
--- a/python/array_record_data_source.py
+++ b/python/array_record_data_source.py
@@ -45,6 +45,7 @@ def __getitem__(self, record_keys: Sequence[int]) -> Sequence[T]:
 from absl import flags
 from absl import logging
 from etils import epath
+from etils import epy
 
 from python import array_record_module