select_device() cleanup (ultralytics#6302)

* `select_device()` cleanup * Update torch_utils.py * Update torch_utils.py * Update torch_utils.py * Update torch_utils.py * Update torch_utils.py
neuralmagic · Jan 15, 2022 · 436ffc4 · 436ffc4
1 parent a1a9c68
commit 436ffc4
Showing 1 changed file with 3 additions and 3 deletions.
diff --git 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()