diff --git a/detect.py b/detect.py
index bf858e4..bc4ce62 100644
--- a/detect.py
+++ b/detect.py
@@ -1,5 +1,7 @@
 import argparse
 
+import torch.backends.cudnn as cudnn
+
 from utils.datasets import *
 from utils.utils import *
 
@@ -36,7 +38,7 @@ def detect(save_img=False):
     vid_path, vid_writer = None, None
     if webcam:
         view_img = True
-        torch.backends.cudnn.benchmark = True  # set True to speed up constant image size inference
+        cudnn.benchmark = True  # set True to speed up constant image size inference
         dataset = LoadStreams(source, img_size=imgsz)
     else:
         save_img = True
diff --git a/train.py b/train.py
index 5163bb1..f545ae1 100644
--- a/train.py
+++ b/train.py
@@ -4,6 +4,7 @@ import torch.distributed as dist
 import torch.nn.functional as F
 import torch.optim as optim
 import torch.optim.lr_scheduler as lr_scheduler
+import torch.utils.data
 from torch.utils.tensorboard import SummaryWriter
 
 import test  # import test.py to get mAP after each epoch