Fix select_device() for Multi-GPU (#6434)

* Fix `select_device()` for Multi-GPU

Possible fix for https://github.com/ultralytics/yolov5/issues/6431

* Update torch_utils.py

* Update torch_utils.py

* Update torch_utils.py

* Update torch_utils.py

* Update

* Update

* Update

* Update

* Update

* Update

* Update

* Update

* Update

utils/torch_utils.py

@@ -53,6 +53,15 @@ def git_describe(path=Path(__file__).parent):  # path must be a directory
         return ''  # not a git repository
 
 
+def device_count():
+    # Returns number of CUDA devices available. Safe version of torch.cuda.device_count().
+    try:
+        cmd = 'nvidia-smi -L | wc -l'
+        return int(subprocess.run(cmd, shell=True, capture_output=True, check=True).stdout.decode().split()[-1])
+    except Exception as e:
+        return 0
+
+
 def select_device(device='', batch_size=0, newline=True):
     # device = 'cpu' or '0' or '0,1,2,3'
     s = f'YOLOv5 🚀 {git_describe() or date_modified()} torch {torch.__version__} '  # string
@@ -61,10 +70,10 @@ 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
-        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'
+        nd = device_count()  # number of CUDA devices
         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)
+        os.environ['CUDA_VISIBLE_DEVICES'] = device  # set environment variable - must be before assert is_available()
+        assert torch.cuda.is_available(), 'CUDA is not available, use `--device cpu` or do not pass a --device'
 
     cuda = not cpu and torch.cuda.is_available()
     if cuda: