Fix device count check (ultralytics#6290)

* Fix device count check()

* Update torch_utils.py

* Update torch_utils.py

* Update hubconf.py

hubconf.py

@@ -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
 
 

utils/torch_utils.py

@@ -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: