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

5 years ago · c099374771
parent b93fcf0708
commit c099374771
1 changed files with 1 additions and 1 deletions
--- 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']