This commit is contained in:
Glenn Jocher 2019-09-19 18:09:16 +02:00
parent de0612ca09
commit 0f3f6c03e7
1 changed files with 1 additions and 0 deletions

View File

@ -22,6 +22,7 @@ def detect(save_txt=False, save_img=False, view_img=False):
model = Darknet(opt.cfg, img_size) model = Darknet(opt.cfg, img_size)
# Load weights # Load weights
attempt_download(weights)
if weights.endswith('.pt'): # pytorch format if weights.endswith('.pt'): # pytorch format
model.load_state_dict(torch.load(weights, map_location=device)['model']) model.load_state_dict(torch.load(weights, map_location=device)['model'])
else: # darknet format else: # darknet format