Always map location to cpu when load checkpoint (#6405)

open-mmlab · Nov 2, 2021 · 599fcc2 · 599fcc2
1 parent db256a1
commit 599fcc2
Showing 1 changed file with 1 addition and 2 deletions.
diff --git a/mmdet/apis/inference.py b/mmdet/apis/inference.py
@@ -39,8 +39,7 @@ def init_detector(config, checkpoint=None, device='cuda:0', cfg_options=None):
     config.model.train_cfg = None
     model = build_detector(config.model, test_cfg=config.get('test_cfg'))
     if checkpoint is not None:
-        map_loc = 'cpu' if device == 'cpu' else None
-        checkpoint = load_checkpoint(model, checkpoint, map_location=map_loc)
+        checkpoint = load_checkpoint(model, checkpoint, map_location='cpu')
         if 'CLASSES' in checkpoint.get('meta', {}):
             model.CLASSES = checkpoint['meta']['CLASSES']
         else: