From c099374771ab386522f98cb777d4fa97a5cf95c2 Mon Sep 17 00:00:00 2001 From: Ross Wightman Date: Sat, 19 Oct 2019 17:27:46 -0700 Subject: [PATCH] Map pretrained checkpoint to cpu to avoid issue with some pretrained checkpoints still having CUDA tensors. Fixes #42 --- timm/models/helpers.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/timm/models/helpers.py b/timm/models/helpers.py index 919bb997..9ac728da 100644 --- a/timm/models/helpers.py +++ b/timm/models/helpers.py @@ -64,7 +64,7 @@ def load_pretrained(model, cfg=None, num_classes=1000, in_chans=3, filter_fn=Non logging.warning("Pretrained model URL is invalid, using random initialization.") return - state_dict = model_zoo.load_url(cfg['url'], progress=False) + state_dict = model_zoo.load_url(cfg['url'], progress=False, map_location='cpu') if in_chans == 1: conv1_name = cfg['first_conv']