FeatureExtractor: avoid loading weights twice

When passing a model_path to the FeatureExtractor, weights are loaded twice.
One the pretrained once and then the given ones. Avoid the unnecessary load.
pull/488/head
Helmut Grohne 2022-01-19 09:25:20 +01:00 committed by Helmut Grohne
parent 6987d99074
commit 42185c4beb
1 changed files with 1 additions and 1 deletions

View File

@ -71,7 +71,7 @@ class FeatureExtractor(object):
model = build_model(
model_name,
num_classes=1,
pretrained=True,
pretrained=not (model_path and check_isfile(model_path)),
use_gpu=device.startswith('cuda')
)
model.eval()