mirrors
/
pytorch-image-models
mirror of https://github.com/huggingface/pytorch-image-models.git
1
0
Fork 0
Code Issues Projects Releases Wiki Activity

Merge pull request #23 from aovoc/patch-1 

Update helpers.py
pull/30/head
Ross Wightman 2019-08-10 15:47:40 -07:00 committed by GitHub
parent 0c874195db 4e7a854dd0
commit f937e2a600
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 2 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
timm/models/helpers.py
View File

@ -7,7 +7,7 @@ from collections import OrderedDict
def load_checkpoint(model, checkpoint_path, use_ema=False):
if checkpoint_path and os.path.isfile(checkpoint_path):
checkpoint = torch.load(checkpoint_path)
checkpoint = torch.load(checkpoint_path, map_location='cpu')
state_dict_key = ''
if isinstance(checkpoint, dict):
state_dict_key = 'state_dict'
@ -32,7 +32,7 @@ def resume_checkpoint(model, checkpoint_path):
optimizer_state = None
resume_epoch = None
if os.path.isfile(checkpoint_path):
checkpoint = torch.load(checkpoint_path)
checkpoint = torch.load(checkpoint_path, map_location='cpu')
if isinstance(checkpoint, dict) and 'state_dict' in checkpoint:
new_state_dict = OrderedDict()
for k, v in checkpoint['state_dict'].items():