diff --git a/yolort/data/coco_eval.py b/yolort/data/coco_eval.py
index ddf0719a..6456d9a4 100644
--- a/yolort/data/coco_eval.py
+++ b/yolort/data/coco_eval.py
@@ -1,4 +1,5 @@
-# Copyright (c) 2021, Zhiqiang Wang. All Rights Reserved.
+# Copyright (c) 2021, yolort team. All rights reserved.
+
 import contextlib
 import copy
 import io
@@ -8,16 +9,16 @@
 
 import numpy as np
 from tabulate import tabulate
-from torchmetrics import Metric
 from torchvision.ops import box_convert
 
 try:
     from pycocotools.coco import COCO
     from pycocotools.cocoeval import COCOeval
+    from torchmetrics import Metric
 except ImportError:
-    COCO, COCOeval = None, None
+    COCO, COCOeval, Metric = None, None, None
 
-from typing import List, Any, Callable, Optional, Union
+from typing import Any, List, Callable, Optional, Union
 
 from ._helper import create_small_table
 from .distributed import all_gather
diff --git a/yolort/data/data_module.py b/yolort/data/data_module.py
index de2bb3bd..ae611d4f 100644
--- a/yolort/data/data_module.py
+++ b/yolort/data/data_module.py
@@ -1,11 +1,16 @@
-# Copyright (c) 2021, Zhiqiang Wang. All Rights Reserved.
+# Copyright (c) 2021, yolort team. All rights reserved.
+
 from pathlib import Path
 from typing import Callable, List, Any, Optional
 
 import torch.utils.data
-from pytorch_lightning import LightningDataModule
 from torch.utils.data.dataset import Dataset
 
+try:
+    from pytorch_lightning import LightningDataModule
+except ImportError:
+    LightningDataModule = None
+
 from .coco import COCODetection
 from .transforms import collate_fn, default_train_transforms, default_val_transforms
 from .voc import VOCDetection