Unverified 提交 e7bf3827 authored 作者: Glenn Jocher's avatar Glenn Jocher 提交者: GitHub

Fix `device` count check (#6290)

* Fix device count check() * Update torch_utils.py * Update torch_utils.py * Update hubconf.py
上级 af001349
......@@ -61,7 +61,7 @@ def _create(name, pretrained=True, channels=3, classes=80, autoshape=True, verbo
except Exception as e:
help_url = 'https://github.com/ultralytics/yolov5/issues/36'
s = 'Cache may be out of date, try `force_reload=True`. See %s for help.' % help_url
s = f'{e}. Cache may be out of date, try `force_reload=True` or see {help_url} for help.'
raise Exception(s) from e
......
......@@ -61,8 +61,9 @@ def select_device(device='', batch_size=0, newline=True):
if cpu:
os.environ['CUDA_VISIBLE_DEVICES'] = '-1' # force torch.cuda.is_available() = False
elif device: # non-cpu device requested
os.environ['CUDA_VISIBLE_DEVICES'] = device # set environment variable
assert torch.cuda.is_available(), f'CUDA unavailable, invalid device {device} requested' # check availability
assert torch.cuda.is_available(), 'CUDA unavailable' # check CUDA is available
assert torch.cuda.device_count() > int(device), f'invalid CUDA device {device} requested' # check index
os.environ['CUDA_VISIBLE_DEVICES'] = device # set environment variable (must be after asserts)
cuda = not cpu and torch.cuda.is_available()
if cuda:
......
Markdown 格式
0%
您添加了 0 到此讨论。请谨慎行事。
请先完成此评论的编辑!
注册 或者 后发表评论