diff --git a/.github/workflows/ci-test.yml b/.github/workflows/ci-test.yml
index 5954a593..297b9ff9 100644
--- a/.github/workflows/ci-test.yml
+++ b/.github/workflows/ci-test.yml
@@ -68,6 +68,9 @@ jobs:
           pip install --user --progress-bar off pycocotools>=2.0.2
           pip install --user --progress-bar off onnx
           pip install --user --progress-bar off onnxruntime
+          pip install --user --progress-bar off pytorch_lightning
+          pip install --user --progress-bar off torchmetrics
+          pip install --user --progress-bar off onnxruntime
           pip install --user onnx_graphsurgeon --index-url https://pypi.ngc.nvidia.com
 
       - name: Lint with flake8
diff --git a/requirements.txt b/requirements.txt
index af9f6bd7..53bd963c 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -24,10 +24,6 @@ ipython
 tabulate
 pandas
 
-# Lightning -----------------------------------
-pytorch_lightning>=1.3.1
-torchmetrics
-
 # extras --------------------------------------
 # pycocotools on PyPI needs python3.7 as minimal
 # pycocotools>=2.0.2  # corresponds to https://github.com/ppwwyyxx/cocoapi
diff --git a/test/test_data_pipeline.py b/test/test_data_pipeline.py
index 4c08e07e..87eb75df 100644
--- a/test/test_data_pipeline.py
+++ b/test/test_data_pipeline.py
@@ -4,7 +4,8 @@
 import numpy as np
 import torch
 from torch import Tensor
-from yolort.data import DetectionDataModule, contains_any_tensor, _helper as data_helper
+from yolort.data import contains_any_tensor, _helper as data_helper
+from yolort.data.data_module import DetectionDataModule
 
 
 def test_contains_any_tensor():
diff --git a/test/test_models_yolov5.py b/test/test_models_yolov5.py
index 93d7a1d6..a1ac3eef 100644
--- a/test/test_models_yolov5.py
+++ b/test/test_models_yolov5.py
@@ -4,7 +4,8 @@
 import torch
 from torch import Tensor
 from torchvision.io import read_image
-from yolort.data import COCOEvaluator, _helper as data_helper
+from yolort.data import _helper as data_helper
+from yolort.data.coco_eval import COCOEvaluator
 from yolort.models import yolov5s
 from yolort.models.transform import YOLOTransform
 from yolort.models.yolo import yolov5_darknet_pan_s_r31
diff --git a/test/test_trainer.py b/test/test_trainer.py
index 5f0844b8..b65ac61c 100644
--- a/test/test_trainer.py
+++ b/test/test_trainer.py
@@ -4,7 +4,8 @@
 
 import pytest
 import pytorch_lightning as pl
-from yolort.data import DetectionDataModule, _helper as data_helper
+from yolort.data import _helper as data_helper
+from yolort.data.data_module import DetectionDataModule
 from yolort.trainer import DefaultTask
 
 
diff --git a/tools/eval_metric.py b/tools/eval_metric.py
index bdda9a88..9bfcb35e 100644
--- a/tools/eval_metric.py
+++ b/tools/eval_metric.py
@@ -7,8 +7,9 @@
 import torch
 import torchvision
 import yolort
-from yolort.data import COCOEvaluator, _helper as data_helper
+from yolort.data import _helper as data_helper
 from yolort.data.coco import COCODetection
+from yolort.data.coco_eval import COCOEvaluator
 from yolort.data.transforms import default_val_transforms, collate_fn
 from yolort.utils.logger import MetricLogger
 
diff --git a/yolort/data/__init__.py b/yolort/data/__init__.py
index 27fd9719..fa9f8ae4 100644
--- a/yolort/data/__init__.py
+++ b/yolort/data/__init__.py
@@ -1,17 +1,5 @@
 # Copyright (c) 2021, yolort team. All rights reserved.
 
 from ._helper import contains_any_tensor
-from .coco_eval import COCOEvaluator
-from .data_module import (
-    DetectionDataModule,
-    VOCDetectionDataModule,
-    COCODetectionDataModule,
-)
 
-__all__ = [
-    "contains_any_tensor",
-    "COCOEvaluator",
-    "DetectionDataModule",
-    "VOCDetectionDataModule",
-    "COCODetectionDataModule",
-]
+__all__ = ["contains_any_tensor"]
diff --git a/yolort/models/_utils.py b/yolort/models/_utils.py
index a11f8305..5d4ee19b 100644
--- a/yolort/models/_utils.py
+++ b/yolort/models/_utils.py
@@ -1,3 +1,5 @@
+# Copyright (c) 2020, yolort team. All rights reserved.
+
 import math
 from typing import Tuple, Optional
 
diff --git a/yolort/models/darknet.py b/yolort/models/darknet.py
index 6c3d9f39..c97b3fda 100644
--- a/yolort/models/darknet.py
+++ b/yolort/models/darknet.py
@@ -1,4 +1,5 @@
 # Copyright (c) 2021, yolort team. All rights reserved.
+
 from .darknetv4 import (
     DarkNetV4,
     darknet_s_r3_1,
diff --git a/yolort/models/darknetv4.py b/yolort/models/darknetv4.py
index d37cbd7d..41d05ad0 100644
--- a/yolort/models/darknetv4.py
+++ b/yolort/models/darknetv4.py
@@ -1,4 +1,5 @@
-# Copyright (c) 2021, Zhiqiang Wang. All Rights Reserved.
+# Copyright (c) 2021, yolort team. All rights reserved.
+
 from typing import Callable, List, Optional, Any
 
 import torch
diff --git a/yolort/models/darknetv6.py b/yolort/models/darknetv6.py
index b4e76c3b..4af755f7 100644
--- a/yolort/models/darknetv6.py
+++ b/yolort/models/darknetv6.py
@@ -1,5 +1,6 @@
 # Copyright (c) 2021, yolort team. All rights reserved.
-from typing import Callable, List, Optional, Any
+
+from typing import Any, Callable, List, Optional
 
 import torch
 from torch import nn, Tensor
diff --git a/yolort/models/path_aggregation_network.py b/yolort/models/path_aggregation_network.py
index faa649d4..fcdf7b35 100644
--- a/yolort/models/path_aggregation_network.py
+++ b/yolort/models/path_aggregation_network.py
@@ -1,5 +1,6 @@
-# Copyright (c) 2021, Zhiqiang Wang. All Rights Reserved.
-from typing import List, Dict, Callable, Optional
+# Copyright (c) 2021, yolort team. All rights reserved.
+
+from typing import Dict, List, Callable, Optional
 
 import torch
 from torch import nn, Tensor
diff --git a/yolort/models/transformer.py b/yolort/models/transformer.py
index fca95ce5..132bc536 100644
--- a/yolort/models/transformer.py
+++ b/yolort/models/transformer.py
@@ -1,9 +1,5 @@
-# Copyright (c) 2021, Zhiqiang Wang. All Rights Reserved.
-"""
-The transformer attention network blocks.
+# Copyright (c) 2021, yolort team. All rights reserved.
 
-Mostly copy-paste from <https://github.com/dingyiwei/yolov5/tree/Transformer>.
-"""
 from typing import Callable, List, Optional
 
 from torch import nn
diff --git a/yolort/trainer/lightning_task.py b/yolort/trainer/lightning_task.py
index 34a8c49a..c8cec1e1 100644
--- a/yolort/trainer/lightning_task.py
+++ b/yolort/trainer/lightning_task.py
@@ -9,7 +9,7 @@
 from pytorch_lightning import LightningModule
 from torch import Tensor
 from torchvision.ops import box_iou
-from yolort.data import COCOEvaluator
+from yolort.data.coco_eval import COCOEvaluator
 
 
 __all__ = ["DefaultTask"]