Skip to content

Commit eac3c03

Browse files
authored
fix device
1 parent 7efd520 commit eac3c03

1 file changed

Lines changed: 1 addition & 1 deletion

File tree

hubconf.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -43,7 +43,7 @@ def _create(name, pretrained=True, channels=3, classes=80, autoshape=True, verbo
4343
device = select_device(('0' if torch.cuda.is_available() else 'cpu') if device is None else device)
4444

4545
if pretrained and channels == 3 and classes == 80:
46-
model = DetectMultiBackend(path, device='cpu') # download/load FP32 model
46+
model = DetectMultiBackend(path, device=device) # download/load FP32 model
4747
else:
4848
cfg = list((Path(__file__).parent / 'models').rglob(f'{path.name}.yaml'))[0] # model.yaml path
4949
model = Model(cfg, channels, classes) # create model

0 commit comments

Comments
 (0)