From 2368603484650e3a18e6c40ec7374beae7bd3391 Mon Sep 17 00:00:00 2001 From: Lornatang Date: Wed, 17 Jun 2020 09:56:26 +0800 Subject: [PATCH] fix refrence bug In torch==1.5, the import of the API has changed. Although it does not interrupt the operation of the program, it seems to me to be an implicit error and may throw an exception in later versions. --- detect.py | 4 +++- train.py | 1 + 2 files changed, 4 insertions(+), 1 deletion(-) 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 7df99c0..6c559c1 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