Unverified Commit 84855b37 authored by Tai-Wang's avatar Tai-Wang Committed by GitHub
Browse files

Fix the memory-leak problem in init_detector (#1045)

parent 5ff7c6ac
......@@ -57,7 +57,7 @@ def init_model(config, checkpoint=None, device='cuda:0'):
config.model.train_cfg = None
model = build_model(config.model, test_cfg=config.get('test_cfg'))
if checkpoint is not None:
checkpoint = load_checkpoint(model, checkpoint)
checkpoint = load_checkpoint(model, checkpoint, map_location='cpu')
if 'CLASSES' in checkpoint['meta']:
model.CLASSES = checkpoint['meta']['CLASSES']
else:
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment