map location to cpu when load checkpoint (#1007)

open-mmlab · Nov 10, 2021 · ec2d4d0 · ec2d4d0
1 parent c0d5e4c
commit ec2d4d0
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/mmpose/apis/inference.py b/mmpose/apis/inference.py
@@ -40,7 +40,7 @@ def init_pose_model(config, checkpoint=None, device='cuda:0'):
     model = build_posenet(config.model)
     if checkpoint is not None:
         # load model checkpoint
-        load_checkpoint(model, checkpoint, map_location=device)
+        load_checkpoint(model, checkpoint, map_location='cpu')
     # save the config in the model for convenience
     model.cfg = config
     model.to(device)