Skip to content
Snippets Groups Projects
Commit e3858425 authored by David de la Iglesia Castro's avatar David de la Iglesia Castro Committed by GitHub
Browse files

Fix cpu inference (#152)


* Add missing map_location

* Add docstring

* Update mmseg/apis/inference.py

Co-authored-by: default avatarJerry Jiarui XU <xvjiarui0826@gmail.com>

* Update inference.py

* Update inference.py

Co-authored-by: default avatarJerry Jiarui XU <xvjiarui0826@gmail.com>
parent 51e4cdef
No related branches found
No related tags found
No related merge requests found
......@@ -16,7 +16,8 @@ def init_segmentor(config, checkpoint=None, device='cuda:0'):
object.
checkpoint (str, optional): Checkpoint path. If left as None, the model
will not load any weights.
device (str, optional) CPU/CUDA device option. Default 'cuda:0'.
Use 'cpu' for loading model on CPU.
Returns:
nn.Module: The constructed segmentor.
"""
......@@ -28,7 +29,7 @@ def init_segmentor(config, checkpoint=None, device='cuda:0'):
config.model.pretrained = None
model = build_segmentor(config.model, test_cfg=config.test_cfg)
if checkpoint is not None:
checkpoint = load_checkpoint(model, checkpoint)
checkpoint = load_checkpoint(model, checkpoint, map_location='cpu')
model.CLASSES = checkpoint['meta']['CLASSES']
model.PALETTE = checkpoint['meta']['PALETTE']
model.cfg = config # save the config in the model for convenience
......
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment