Skip to content

GitLab

Unverified Commit 3fd5ea26 authored by aldakata's avatar aldakata Committed by GitHub
Browse files

Explicitly setting torch.cuda.device at init_model (#1056)

parent c06c0189
Hide whitespace changes
Inline Side-by-side
Showing with 1 addition and 0 deletions
mmdet3d/apis/inference.py
View file @ 3fd5ea26
...@@ -65,6 +65,7 @@ def init_model(config, checkpoint=None, device='cuda:0'): ...@@ -65,6 +65,7 @@ def init_model(config, checkpoint=None, device='cuda:0'):
if 'PALETTE' in checkpoint['meta']: # 3D Segmentor if 'PALETTE' in checkpoint['meta']: # 3D Segmentor
model.PALETTE = checkpoint['meta']['PALETTE'] model.PALETTE = checkpoint['meta']['PALETTE']
model.cfg = config # save the config in the model for convenience model.cfg = config # save the config in the model for convenience
torch.cuda.set_device(device)
model.to(device) model.to(device)
model.eval() model.eval()
return model return model
......
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