updates
This commit is contained in:
parent
521b4c02ff
commit
1c72eb03f0
2
test.py
2
test.py
|
@ -31,7 +31,7 @@ num_classes = int(data_config['classes'])
|
||||||
model = Darknet(opt.cfg, opt.img_size)
|
model = Darknet(opt.cfg, opt.img_size)
|
||||||
|
|
||||||
# Load weights
|
# Load weights
|
||||||
if opt.weights_path('.weights'): # darknet format
|
if opt.weights_path.endswith('.weights'): # darknet format
|
||||||
load_weights(model, opt.weights_path)
|
load_weights(model, opt.weights_path)
|
||||||
elif opt.weights_path.endswith('.pt'): # pytorch format
|
elif opt.weights_path.endswith('.pt'): # pytorch format
|
||||||
checkpoint = torch.load(opt.weights_path, map_location='cpu')
|
checkpoint = torch.load(opt.weights_path, map_location='cpu')
|
||||||
|
|
Loading…
Reference in New Issue