Commit 00b15d32 authored by Anjith GEORGE's avatar Anjith GEORGE

small fix

parent 27898ca9
Pipeline #28688 failed with stage
in 19 minutes and 27 seconds
......@@ -59,22 +59,16 @@ class MCDeepPixBiSExtractor(Extractor):
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('Failed to load the model file : {}'.format(model_file))
raise ValueError('Failed to load the model file : {}'.format(model_file))
if 'state_dict' in cp:
self.network.load_state_dict(cp['state_dict'])
elif: ## check this part
self.network.load_state_dict(cp)
else:
raise ValueError('Failed to load the state_dict for model file: {}'.format(model_file))
......
Markdown is supported
0%
or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment