Merge pull request #107 from Lornatang/fix-reference-bugs

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

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