From 12065f3bc5dbe7abb02a3412b640cfffafdbf06f Mon Sep 17 00:00:00 2001
From: zhiqiang <zhiqwang@outlook.com>
Date: Thu, 18 Mar 2021 23:27:29 +0800
Subject: [PATCH] Fixing module loading

---
 test/test_models.py          |  4 ++--
 yolort/models/transformer.py |  2 +-
 yolort/models/yolo.py        | 10 +++++-----
 3 files changed, 8 insertions(+), 8 deletions(-)

diff --git a/test/test_models.py b/test/test_models.py
index ecd24039..462787ca 100644
--- a/test/test_models.py
+++ b/test/test_models.py
@@ -2,7 +2,7 @@
 import torch
 
 from yolort.models.backbone_utils import darknet_pan_backbone
-from yolort.models.yolotr import darknet_pan_tr_backbone
+from yolort.models.transformer import darknet_tan_backbone
 from yolort.models.anchor_utils import AnchorGenerator
 from yolort.models.box_head import YoloHead, PostProcess, SetCriterion
 
@@ -112,7 +112,7 @@ def _init_test_backbone_with_pan_tr(self):
         backbone_name = 'darknet_s_r4_0'
         depth_multiple = 0.33
         width_multiple = 0.5
-        backbone_with_fpn_tr = darknet_pan_tr_backbone(backbone_name, depth_multiple, width_multiple)
+        backbone_with_fpn_tr = darknet_tan_backbone(backbone_name, depth_multiple, width_multiple)
         return backbone_with_fpn_tr
 
     def test_backbone_with_pan_tr(self):
diff --git a/yolort/models/transformer.py b/yolort/models/transformer.py
index dbeabe2b..f97a6705 100644
--- a/yolort/models/transformer.py
+++ b/yolort/models/transformer.py
@@ -24,7 +24,7 @@ def darknet_tan_backbone(
     version: str = 'v4.0',
 ):
     """
-    Constructs a specified DarkNet backbone with PAN on top. Freezes the specified number of
+    Constructs a specified DarkNet backbone with TAN on top. Freezes the specified number of
     layers in the backbone.
 
     Examples::
diff --git a/yolort/models/yolo.py b/yolort/models/yolo.py
index 43c5a7a8..3b7bb95f 100644
--- a/yolort/models/yolo.py
+++ b/yolort/models/yolo.py
@@ -7,7 +7,7 @@
 from torchvision.models.utils import load_state_dict_from_url
 
 from .backbone_utils import darknet_pan_backbone
-from .transformer import darknet_pan_tr_backbone
+from .transformer import darknet_tan_backbone
 from .anchor_utils import AnchorGenerator
 from .box_head import YoloHead, SetCriterion, PostProcess
 
@@ -15,7 +15,7 @@
 
 __all__ = ['YOLO', 'yolov5_darknet_pan_s_r31', 'yolov5_darknet_pan_m_r31', 'yolov5_darknet_pan_l_r31',
            'yolov5_darknet_pan_s_r40', 'yolov5_darknet_pan_m_r40', 'yolov5_darknet_pan_l_r40',
-           'yolov5_darknet_pan_s_tr']
+           'yolov5_darknet_tan_s_r40']
 
 
 class YOLO(nn.Module):
@@ -303,8 +303,8 @@ def yolov5_darknet_pan_l_r40(pretrained: bool = False, progress: bool = True, nu
                                pretrained=pretrained, progress=progress, num_classes=num_classes, **kwargs)
 
 
-def yolov5_darknet_pan_s_tr(pretrained: bool = False, progress: bool = True, num_classes: int = 80,
-                            **kwargs: Any) -> YOLO:
+def yolov5_darknet_tan_s_r40(pretrained: bool = False, progress: bool = True, num_classes: int = 80,
+                             **kwargs: Any) -> YOLO:
     r"""yolov5 small with a transformer block model from
     `"dingyiwei/yolov5" <https://github.com/ultralytics/yolov5/pull/2333>`_.
     Args:
@@ -317,7 +317,7 @@ def yolov5_darknet_pan_s_tr(pretrained: bool = False, progress: bool = True, num
     width_multiple = 0.5
     version = 'v4.0'
 
-    backbone = darknet_pan_tr_backbone(backbone_name, depth_multiple, width_multiple, version=version)
+    backbone = darknet_tan_backbone(backbone_name, depth_multiple, width_multiple, version=version)
 
     anchor_grids = [[10, 13, 16, 30, 33, 23],
                     [30, 61, 62, 45, 59, 119],