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

ZwwWayne · Jul 18, 2022 · e0c15dd · e0c15dd
1 parent e8e04d2
commit e0c15dd
Showing 1 changed file with 1 addition and 2 deletions.
diff --git a/mmdet/apis/inference.py b/mmdet/apis/inference.py
@@ -38,8 +38,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: