From 08f051c1d42c71836b672dfcf7c0c57b4c63859e Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Fri, 8 Feb 2019 23:20:41 +0100 Subject: [PATCH] updates --- detect.py | 4 +--- test.py | 4 +--- 2 files changed, 2 insertions(+), 6 deletions(-) diff --git a/detect.py b/detect.py index e7070b56..5fbe3299 100755 --- a/detect.py +++ b/detect.py @@ -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) diff --git a/test.py b/test.py index bcb54f99..8b429660 100644 --- a/test.py +++ b/test.py @@ -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)