This commit is contained in:
Glenn Jocher 2019-02-08 23:20:41 +01:00
parent 8dec060504
commit 08f051c1d4
2 changed files with 2 additions and 6 deletions

View File

@ -30,9 +30,7 @@ def detect(
if weights.endswith('.pt'): # pytorch format
if weights.endswith('weights/yolov3.pt') and not os.path.isfile(weights):
os.system('wget https://storage.googleapis.com/ultralytics/yolov3.pt -O ' + weights)
checkpoint = torch.load(weights, map_location='cpu')
model.load_state_dict(checkpoint['model'])
del checkpoint
model.load_state_dict(torch.load(weights, map_location='cpu')['model'])
else: # darknet format
load_darknet_weights(model, weights)

View File

@ -29,9 +29,7 @@ def test(
# Load weights
if weights.endswith('.pt'): # pytorch format
checkpoint = torch.load(weights, map_location='cpu')
model.load_state_dict(checkpoint['model'])
del checkpoint
model.load_state_dict(torch.load(weights, map_location='cpu')['model'])
else: # darknet format
load_darknet_weights(model, weights)