xclele/yolov10
1
0
Fork 0
mirror of https://github.com/THU-MIG/yolov10.git synced 2025-05-23 13:34:23 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

Fix device counting method to account for double-digit device IDs (#8502)

Browse Source 
Co-authored-by: Glenn Jocher <glenn.jocher@ultralytics.com>
This commit is contained in:
Omar Duhaiby 2024-03-01 12:19:06 +01:00 committed by GitHub
parent 59ed47c448
commit 3c6b9b6688
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 1 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
ultralytics/utils/torch_utils.py
View File

@ -115,7 +115,7 @@ def select_device(device="", batch=0, newline=False, verbose=True):
device = "0" device = "0"
visible = os.environ.get("CUDA_VISIBLE_DEVICES", None) visible = os.environ.get("CUDA_VISIBLE_DEVICES", None)
os.environ["CUDA_VISIBLE_DEVICES"] = device # set environment variable - must be before assert is_available() 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) LOGGER.info(s)
install = ( install = (
"See https://pytorch.org/get-started/locally/ for up-to-date torch install instructions if no " "See https://pytorch.org/get-started/locally/ for up-to-date torch install instructions if no "