From eb6601b0cf9d37dff879b819d542297e21201287 Mon Sep 17 00:00:00 2001 From: Henry Date: Fri, 14 Jan 2022 17:57:01 +0100 Subject: [PATCH 1/2] Fixing bug multi-gpu training This solves this issue: https://github.com/ultralytics/yolov5/issues/6297#issue-1103853348 --- utils/torch_utils.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/utils/torch_utils.py b/utils/torch_utils.py index 060768e8251b..ede653f106d8 100644 --- a/utils/torch_utils.py +++ b/utils/torch_utils.py @@ -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 elif device: # non-cpu device requested 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) cuda = not cpu and torch.cuda.is_available() From 9ee92ed0cc786a15a10252277ec4cd43ff4bb7c9 Mon Sep 17 00:00:00 2001 From: Henry Date: Fri, 14 Jan 2022 18:25:29 +0100 Subject: [PATCH 2/2] Update torch_utils.py for pep8 --- utils/torch_utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/utils/torch_utils.py b/utils/torch_utils.py index ede653f106d8..451bcdd29b7c 100644 --- a/utils/torch_utils.py +++ b/utils/torch_utils.py @@ -63,7 +63,7 @@ def select_device(device='', batch_size=0, newline=True): elif device: # non-cpu device requested assert torch.cuda.is_available(), 'CUDA unavailable' # check CUDA is available 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 + 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) cuda = not cpu and torch.cuda.is_available()