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/datasets.py

@@ -29,13 +29,13 @@ from tqdm import tqdm
 from utils.augmentations import Albumentations, augment_hsv, copy_paste, letterbox, mixup, random_perspective
 from utils.general import (LOGGER, NUM_THREADS, check_dataset, check_requirements, check_yaml, clean_str,
                            segments2boxes, xyn2xy, xywh2xyxy, xywhn2xyxy, xyxy2xywhn)
-from utils.torch_utils import torch_distributed_zero_first
+from utils.torch_utils import device_count, torch_distributed_zero_first
 
 # Parameters
 HELP_URL = 'https://github.com/ultralytics/yolov5/wiki/Train-Custom-Data'
 IMG_FORMATS = ['bmp', 'dng', 'jpeg', 'jpg', 'mpo', 'png', 'tif', 'tiff', 'webp']  # include image suffixes
 VID_FORMATS = ['asf', 'avi', 'gif', 'm4v', 'mkv', 'mov', 'mp4', 'mpeg', 'mpg', 'wmv']  # include video suffixes
-DEVICE_COUNT = max(torch.cuda.device_count(), 1)
+DEVICE_COUNT = max(device_count(), 1)  # number of CUDA devices
 
 # Get orientation exif tag
 for orientation in ExifTags.TAGS.keys():