Map pretrained checkpoint to cpu to avoid issue with some pretrained checkpoints still having CUDA tensors. Fixes #42

Author: rwightman

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']