From 688847d1f329db28badd4f8db851be5792cadcf9 Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Fri, 14 Jan 2022 15:48:15 -1000 Subject: [PATCH] `select_device()` cleanup (#6302) * `select_device()` cleanup * Update torch_utils.py * Update torch_utils.py * Update torch_utils.py * Update torch_utils.py * Update torch_utils.py --- utils/torch_utils.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/utils/torch_utils.py b/utils/torch_utils.py index 451bcdd29b7c..7e464190f9ba 100644 --- a/utils/torch_utils.py +++ b/utils/torch_utils.py @@ -61,9 +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 - 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 + nd = torch.cuda.device_count() # number of CUDA devices + assert torch.cuda.is_available(), 'CUDA is not available, use `--device cpu` or do not pass a --device' + assert nd > int(max(device.split(','))), f'Invalid `--device {device}` request, valid devices are 0 - {nd - 1}' os.environ['CUDA_VISIBLE_DEVICES'] = device # set environment variable (must be after asserts) cuda = not cpu and torch.cuda.is_available()