diff --git a/utils.py b/picotron/utils.py
similarity index 100%
rename from utils.py
rename to picotron/utils.py
diff --git a/train.py b/train.py
index 593ffda..a7898e3 100644
--- a/train.py
+++ b/train.py
@@ -22,7 +22,7 @@ from transformers import AutoConfig
 import numpy as np
 from picotron.parallel.tensor_parallel.tensor_parallel import TensorParallel
 import picotron.process_group_manager as pgm
-from utils import MicroBatchDataLoader, set_all_seed, print, to_readable_format, save_checkpoint, load_checkpoint
+from picotron.parallel.utils import MicroBatchDataLoader, set_all_seed, print, to_readable_format, save_checkpoint, load_checkpoint
 from picotron.process_group_manager import setup_process_group_manager
 from picotron.parallel.pipeline_parallel import train_step_pipeline_1f1b, train_step_pipeline_afab, PipelineParallel
 from picotron.parallel.data_parallel.data_parallel_bucket import DataParallel