diff --git a/yolort/models/_utils.py b/yolort/models/_utils.py
index 3ff332ccc..e646d48ba 100644
--- a/yolort/models/_utils.py
+++ b/yolort/models/_utils.py
@@ -5,7 +5,7 @@
 from torch import nn, Tensor
 from torchvision.ops import box_convert, box_iou
 
-from yolort.utils import ModuleStateUpdate
+from yolort.utils.update_module_state import ModuleStateUpdate
 from yolort.v5 import load_yolov5_model, get_yolov5_size
 
 
diff --git a/yolort/utils/update_module_state.py b/yolort/utils/update_module_state.py
index 34214bab8..e12be8fdc 100644
--- a/yolort/utils/update_module_state.py
+++ b/yolort/utils/update_module_state.py
@@ -4,7 +4,7 @@
 
 from torch import nn
 
-from yolort.models import yolo
+from yolort.models.yolo import build_model
 from yolort.v5 import get_yolov5_size
 
 
@@ -56,7 +56,7 @@ def __init__(
         weights_name = (
             f"yolov5_darknet_pan_{yolov5_size}_{version.replace('.', '')}_coco"
         )
-        model = yolo.build_model(
+        model = build_model(
             backbone_name,
             depth_multiple,
             width_multiple,