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

pull/88/head
Ross Wightman 5 years ago
parent b9496995d8
commit f098fda2ca

@ -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()

Loading…
Cancel
Save