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.
pull/1/head
Lornatang 5 years ago
parent a557b7d924
commit 2368603484

@ -1,5 +1,7 @@
import argparse import argparse
import torch.backends.cudnn as cudnn
from utils.datasets import * from utils.datasets import *
from utils.utils import * from utils.utils import *
@ -36,7 +38,7 @@ def detect(save_img=False):
vid_path, vid_writer = None, None vid_path, vid_writer = None, None
if webcam: if webcam:
view_img = True 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) dataset = LoadStreams(source, img_size=imgsz)
else: else:
save_img = True save_img = True

@ -4,6 +4,7 @@ import torch.distributed as dist
import torch.nn.functional as F import torch.nn.functional as F
import torch.optim as optim import torch.optim as optim
import torch.optim.lr_scheduler as lr_scheduler import torch.optim.lr_scheduler as lr_scheduler
import torch.utils.data
from torch.utils.tensorboard import SummaryWriter from torch.utils.tensorboard import SummaryWriter
import test # import test.py to get mAP after each epoch import test # import test.py to get mAP after each epoch

Loading…
Cancel
Save