Skip to content
Snippets Groups Projects

Deeppixbis extractor

Merged Guillaume HEUSCH requested to merge deeppixbis-extractor into master
1 unresolved thread
1 file
+ 4
11
Compare changes
  • Side-by-side
  • Inline
@@ -52,24 +52,17 @@ class DeepPixBiSExtractor(Extractor):
logger.debug("No pretrained file provided")
pass
else:
# With the new training
logger.debug('Starting to load the pretrained PAD model')
try:
cp = torch.load(model_file)
except:
try:
cp= torch.load(model_file,map_location=lambda storage,loc:storage)
except:
raise ValueError('Could not load the model')
raise ValueError('Failed to load the model file : {}'.format(model_file))
if 'state_dict' in cp:
self.network.load_state_dict(cp['state_dict'])
else: ## check this part
self.network.load_state_dict(cp)
else:
raise ValueError('Failed to load the state_dict for model file: {}'.format(model_file))
logger.debug('Loaded the pretrained PAD model')
self.network.eval()
Loading