From 4e7a854dd07db4453539f956f998a3fda306eb50 Mon Sep 17 00:00:00 2001
From: Minqin Chen <21521090@zju.edu.cn>
Date: Sun, 11 Aug 2019 04:21:39 +0800
Subject: [PATCH] Update helpers.py

Fixing out of memory error by loading the checkpoint onto the CPU.
---
 timm/models/helpers.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/timm/models/helpers.py b/timm/models/helpers.py
index 1deff273..a37f5062 100644
--- a/timm/models/helpers.py
+++ b/timm/models/helpers.py
@@ -7,7 +7,7 @@ from collections import OrderedDict
 
 def load_checkpoint(model, checkpoint_path, use_ema=False):
     if checkpoint_path and os.path.isfile(checkpoint_path):
-        checkpoint = torch.load(checkpoint_path)
+        checkpoint = torch.load(checkpoint_path, map_location='cpu')
         state_dict_key = ''
         if isinstance(checkpoint, dict):
             state_dict_key = 'state_dict'
@@ -32,7 +32,7 @@ def resume_checkpoint(model, checkpoint_path):
     optimizer_state = None
     resume_epoch = None
     if os.path.isfile(checkpoint_path):
-        checkpoint = torch.load(checkpoint_path)
+        checkpoint = torch.load(checkpoint_path, map_location='cpu')
         if isinstance(checkpoint, dict) and 'state_dict' in checkpoint:
             new_state_dict = OrderedDict()
             for k, v in checkpoint['state_dict'].items():