From f098fda2ca48bf66b611237787b52b9614003a8b Mon Sep 17 00:00:00 2001 From: Ross Wightman Date: Wed, 12 Feb 2020 13:23:56 -0800 Subject: [PATCH] Add map_location='cpu' to ModelEma resume, should improve #72 --- timm/utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/timm/utils.py b/timm/utils.py index 59d2bcd0..1da69a96 100644 --- a/timm/utils.py +++ b/timm/utils.py @@ -259,7 +259,7 @@ class ModelEma: p.requires_grad_(False) def _load_checkpoint(self, checkpoint_path): - checkpoint = torch.load(checkpoint_path) + checkpoint = torch.load(checkpoint_path, map_location='cpu') assert isinstance(checkpoint, dict) if 'state_dict_ema' in checkpoint: new_state_dict = OrderedDict()