diff --git a/realesrgan/train.py b/realesrgan/train.py
index d985a79..e9ad4dd 100644
--- a/realesrgan/train.py
+++ b/realesrgan/train.py
@@ -1,9 +1,10 @@
+# flake8: noqa
 import os.path as osp
 from basicsr.train import train_pipeline
 
-import realesrgan.archs  # noqa: F401
-import realesrgan.data  # noqa: F401
-import realesrgan.models  # noqa: F401
+from .archs import *
+from .data import *
+from .models import *
 
 if __name__ == '__main__':
     root_path = osp.abspath(osp.join(__file__, osp.pardir))
diff --git a/setup.cfg b/setup.cfg
index 4eb6529..2293ad7 100644
--- a/setup.cfg
+++ b/setup.cfg
@@ -16,7 +16,7 @@ split_before_expression_after_opening_paren = true
 line_length = 120
 multi_line_output = 0
 known_standard_library = pkg_resources,setuptools
-known_first_party = basicsr  # modify it!
+known_first_party = realesrgan
 known_third_party = basicsr,cv2,numpy,torch
 no_lines_before = STDLIB,LOCALFOLDER
 default_section = THIRDPARTY