From 3c6b9b66883d0e84ddb2585f06b665e4f6ca585e Mon Sep 17 00:00:00 2001 From: Omar Duhaiby <3omarz@gmail.com> Date: Fri, 1 Mar 2024 12:19:06 +0100 Subject: [PATCH] Fix device counting method to account for double-digit device IDs (#8502) Co-authored-by: Glenn Jocher --- ultralytics/utils/torch_utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ultralytics/utils/torch_utils.py b/ultralytics/utils/torch_utils.py index ed5bc2ed..bf34b1d6 100644 --- a/ultralytics/utils/torch_utils.py +++ b/ultralytics/utils/torch_utils.py @@ -115,7 +115,7 @@ def select_device(device="", batch=0, newline=False, verbose=True): device = "0" visible = os.environ.get("CUDA_VISIBLE_DEVICES", None) os.environ["CUDA_VISIBLE_DEVICES"] = device # set environment variable - must be before assert is_available() - if not (torch.cuda.is_available() and torch.cuda.device_count() >= len(device.replace(",", ""))): + if not (torch.cuda.is_available() and torch.cuda.device_count() >= len(device.split(","))): LOGGER.info(s) install = ( "See https://pytorch.org/get-started/locally/ for up-to-date torch install instructions if no "