diff --git a/src/cerebras/modelzoo/data_preparation/vision/dit/create_imagenet_latents.py b/src/cerebras/modelzoo/data_preparation/vision/dit/create_imagenet_latents.py
index 3d8d17fd..f64928a7 100644
--- a/src/cerebras/modelzoo/data_preparation/vision/dit/create_imagenet_latents.py
+++ b/src/cerebras/modelzoo/data_preparation/vision/dit/create_imagenet_latents.py
@@ -13,6 +13,7 @@
 # limitations under the License.
 
 # isort: off
+import os
 import sys
 
 sys.path.append(os.path.join(os.path.dirname(__file__), "../../../"))
@@ -22,7 +23,6 @@
 import glob
 import json
 import logging
-import os
 import re
 import shutil
 from datetime import datetime