Unverified 提交 a1a9c688 authored 作者: Henry's avatar Henry 提交者: GitHub

Fixing bug multi-gpu training (#6299)

* Fixing bug multi-gpu training This solves this issue: https://github.com/ultralytics/yolov5/issues/6297#issue-1103853348 * Update torch_utils.py for pep8
上级 e7bf3827
...@@ -62,7 +62,8 @@ def select_device(device='', batch_size=0, newline=True): ...@@ -62,7 +62,8 @@ def select_device(device='', batch_size=0, newline=True):
os.environ['CUDA_VISIBLE_DEVICES'] = '-1' # force torch.cuda.is_available() = False os.environ['CUDA_VISIBLE_DEVICES'] = '-1' # force torch.cuda.is_available() = False
elif device: # non-cpu device requested elif device: # non-cpu device requested
assert torch.cuda.is_available(), 'CUDA unavailable' # check CUDA is available 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 device_list = [int(val) for val in device.replace(',', '')]
assert all([torch.cuda.device_count() > element for element in device_list]), f'invalid CUDA device {device} requested' # check index
os.environ['CUDA_VISIBLE_DEVICES'] = device # set environment variable (must be after asserts) os.environ['CUDA_VISIBLE_DEVICES'] = device # set environment variable (must be after asserts)
cuda = not cpu and torch.cuda.is_available() cuda = not cpu and torch.cuda.is_available()
......
Markdown 格式
0%
您添加了 0 到此讨论。请谨慎行事。
请先完成此评论的编辑!
注册 或者 后发表评论