mirrors/pytorch-image-models
1
0
Fork 0
mirror of https://github.com/huggingface/pytorch-image-models.git synced 2025-06-03 15:01:08 +08:00
Code Issues Projects Releases Wiki Activity

Add map_location='cpu' to ModelEma resume, should improve #72

Browse Source
This commit is contained in:
Ross Wightman 2020-02-12 13:23:56 -08:00
parent b9496995d8
commit f098fda2ca
1 changed files with 1 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
timm/utils.py
View File

@ -259,7 +259,7 @@ class ModelEma:
p.requires_grad_(False) p.requires_grad_(False)
def _load_checkpoint(self, checkpoint_path): def _load_checkpoint(self, checkpoint_path):
checkpoint = torch.load(checkpoint_path) checkpoint = torch.load(checkpoint_path, map_location='cpu')
assert isinstance(checkpoint, dict) assert isinstance(checkpoint, dict)
if 'state_dict_ema' in checkpoint: if 'state_dict_ema' in checkpoint:
new_state_dict = OrderedDict() new_state_dict = OrderedDict()