Skip to content

GitLab

"...kompute/git@developer.sourcefind.cn:OpenDAS/ollama.git" did not exist on "97b02a8981110c76d901e8b5f96af514ee0326f3"
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'):
if 'PALETTE' in checkpoint['meta']: # 3D Segmentor
model.PALETTE = checkpoint['meta']['PALETTE']
model.cfg = config # save the config in the model for convenience
torch.cuda.set_device(device)
model.to(device)
model.eval()
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