Skip to content

Commit f937e2a

Browse files
authored
Merge pull request #23 from aovoc/patch-1
Update helpers.py
2 parents 0c87419 + 4e7a854 commit f937e2a

File tree

1 file changed

+2
-2
lines changed

1 file changed

+2
-2
lines changed

timm/models/helpers.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,7 @@
77

88
def load_checkpoint(model, checkpoint_path, use_ema=False):
99
if checkpoint_path and os.path.isfile(checkpoint_path):
10-
checkpoint = torch.load(checkpoint_path)
10+
checkpoint = torch.load(checkpoint_path, map_location='cpu')
1111
state_dict_key = ''
1212
if isinstance(checkpoint, dict):
1313
state_dict_key = 'state_dict'
@@ -32,7 +32,7 @@ def resume_checkpoint(model, checkpoint_path):
3232
optimizer_state = None
3333
resume_epoch = None
3434
if os.path.isfile(checkpoint_path):
35-
checkpoint = torch.load(checkpoint_path)
35+
checkpoint = torch.load(checkpoint_path, map_location='cpu')
3636
if isinstance(checkpoint, dict) and 'state_dict' in checkpoint:
3737
new_state_dict = OrderedDict()
3838
for k, v in checkpoint['state_dict'].items():

0 commit comments

Comments
 (0)