diff --git a/llm/data/utils.py b/llm/data/utils.py
index 85de2de..1d62090 100644
--- a/llm/data/utils.py
+++ b/llm/data/utils.py
@@ -1,6 +1,8 @@
 """data utils"""
 
 import json
+import logging
+from torch.distributed import get_rank, is_initialized
 
 
 def read_text(path, debug=None):