From a5ceef8ca33f6c38af2e5f767987c253cb74d5c6 Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Thu, 13 Jan 2022 21:23:03 -1000 Subject: [PATCH] Fix `device` count check (#6290) * Fix device count check() * Update torch_utils.py * Update torch_utils.py * Update hubconf.py --- hubconf.py | 2 +- utils/torch_utils.py | 5 +++-- 2 files changed, 4 insertions(+), 3 deletions(-) diff --git a/hubconf.py b/hubconf.py index 6bf4b0b0265f..55d15abe2ac5 100644 --- a/hubconf.py +++ b/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 diff --git a/utils/torch_utils.py b/utils/torch_utils.py index cddb173948fb..060768e8251b 100644 --- a/utils/torch_utils.py +++ b/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: