Skip to content
Snippets Groups Projects
Commit 00b15d32 authored by Anjith GEORGE's avatar Anjith GEORGE
Browse files

small fix

parent 27898ca9
No related branches found
No related tags found
1 merge request!30Mc deep pixbis
Pipeline #28688 failed
...@@ -59,22 +59,16 @@ class MCDeepPixBiSExtractor(Extractor): ...@@ -59,22 +59,16 @@ class MCDeepPixBiSExtractor(Extractor):
pass pass
else: else:
# With the new training # With the new training
logger.debug('Starting to load the pretrained PAD model') logger.debug('Starting to load the pretrained PAD model')
try: try:
cp = torch.load(model_file) cp = torch.load(model_file)
except: except:
try: raise ValueError('Failed to load the model file : {}'.format(model_file))
cp= torch.load(model_file,map_location=lambda storage,loc:storage)
except:
raise ValueError('Failed to load the model file : {}'.format(model_file))
if 'state_dict' in cp: if 'state_dict' in cp:
self.network.load_state_dict(cp['state_dict']) self.network.load_state_dict(cp['state_dict'])
elif: ## check this part
self.network.load_state_dict(cp)
else: else:
raise ValueError('Failed to load the state_dict for model file: {}'.format(model_file)) raise ValueError('Failed to load the state_dict for model file: {}'.format(model_file))
......
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Please register or to comment