From 94c89bb509156a13acd493e945fa9f42b5394771 Mon Sep 17 00:00:00 2001
From: Zhiqiang Wang <zhiqwang@outlook.com>
Date: Tue, 21 Sep 2021 02:59:51 -0400
Subject: [PATCH] Rename load_model to load_yolov5_model

---
 yolort/utils/update_module_state.py | 7 +++----
 yolort/v5/__init__.py               | 2 +-
 yolort/v5/helper.py                 | 4 ++--
 yolort/v5/utils/general.py          | 3 ++-
 4 files changed, 8 insertions(+), 8 deletions(-)

diff --git a/yolort/utils/update_module_state.py b/yolort/utils/update_module_state.py
index ffddc7740..ed795ba61 100644
--- a/yolort/utils/update_module_state.py
+++ b/yolort/utils/update_module_state.py
@@ -5,6 +5,7 @@
 from torch import nn
 
 from yolort.models import yolo
+from yolort.v5 import load_yolov5_model
 
 
 ARCHITECTURE_MAPS = {
@@ -41,9 +42,7 @@ def update_module_state_from_ultralytics(
             Default: True.
         verbose (bool): print all information to screen. Default: True.
     """
-
-    from yolort.v5 import load_model
-    model = load_model(model_path, autoshape=False, verbose=verbose)
+    model = load_yolov5_model(model_path, autoshape=False, verbose=verbose)
 
     key_arch = f'{arch}_{feature_fusion_type.lower()}_v4.0'
     if key_arch not in ARCHITECTURE_MAPS:
@@ -141,7 +140,7 @@ def attach_parameters_heads(self, state_dict, name):
 def rgetattr(obj, attr, *args):
     """
     Nested version of getattr.
-    See <https://stackoverflow.com/questions/31174295/getattr-and-setattr-on-nested-objects>
+    Ref: https://stackoverflow.com/questions/31174295/getattr-and-setattr-on-nested-objects
     """
     def _getattr(obj, attr):
         return getattr(obj, attr, *args)
diff --git a/yolort/v5/__init__.py b/yolort/v5/__init__.py
index e5e1a3d31..8052d2c91 100644
--- a/yolort/v5/__init__.py
+++ b/yolort/v5/__init__.py
@@ -1,3 +1,3 @@
 from .models import *
 from .utils import *
-from .helper import add_yolov5_context, load_model
+from .helper import add_yolov5_context, load_yolov5_model
diff --git a/yolort/v5/helper.py b/yolort/v5/helper.py
index 737d09867..4ed85c573 100644
--- a/yolort/v5/helper.py
+++ b/yolort/v5/helper.py
@@ -7,7 +7,7 @@
 from .models.yolo import Model
 from .utils import attempt_download, set_logging
 
-__all__ = ['add_yolov5_context', 'load_model']
+__all__ = ['add_yolov5_context', 'load_yolov5_model']
 
 
 @contextlib.contextmanager
@@ -27,7 +27,7 @@ def add_yolov5_context():
         sys.path.remove(path_ultralytics_yolov5)
 
 
-def load_model(model_path: str, autoshape: bool = False, verbose: bool = True):
+def load_yolov5_model(model_path: str, autoshape: bool = False, verbose: bool = True):
     """
     Creates a specified YOLOv5 model
 
diff --git a/yolort/v5/utils/general.py b/yolort/v5/utils/general.py
index d2fc0f99e..6818b5d78 100644
--- a/yolort/v5/utils/general.py
+++ b/yolort/v5/utils/general.py
@@ -87,7 +87,8 @@ def methods(instance):
 def set_logging(rank=-1, verbose=True):
     logging.basicConfig(
         format="%(message)s",
-        level=logging.INFO if (verbose and rank in [-1, 0]) else logging.WARN)
+        level=logging.INFO if (verbose and rank in [-1, 0]) else logging.WARN,
+    )
 
 
 def init_seeds(seed=0):