diff --git a/lib/datasets/test_utils.py b/lib/datasets/test_utils.py
index ca30512..2245f87 100644
--- a/lib/datasets/test_utils.py
+++ b/lib/datasets/test_utils.py
@@ -1,6 +1,9 @@
 ##################################################
 # Copyright (c) Xuanyi Dong [GitHub D-X-Y], 2019 #
 ##################################################
+import os
+
+
 def test_imagenet_data(imagenet):
   total_length = len(imagenet)
   assert total_length == 1281166 or total_length == 50000, 'The length of ImageNet is wrong : {}'.format(total_length)