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

5 years ago · f098fda2ca
parent b9496995d8
commit f098fda2ca
1 changed files with 1 additions and 1 deletions
--- 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()