From 8e60fc9276c80071c84da953ff7ec11fe7eaad58 Mon Sep 17 00:00:00 2001
From: Glenn Jocher <glenn.jocher@ultralytics.com>
Date: Sun, 11 Jun 2023 17:58:10 +0200
Subject: [PATCH] Add `device=cuda` support (#3133)
---
ultralytics/yolo/utils/torch_utils.py | 2 ++
1 file changed, 2 insertions(+)
diff --git a/ultralytics/yolo/utils/torch_utils.py b/ultralytics/yolo/utils/torch_utils.py
index f576ccd4..cb3725a3 100644
--- a/ultralytics/yolo/utils/torch_utils.py
+++ b/ultralytics/yolo/utils/torch_utils.py
@@ -64,6 +64,8 @@ def select_device(device='', batch=0, newline=False, verbose=True):
if cpu or mps:
os.environ['CUDA_VISIBLE_DEVICES'] = '-1' # force torch.cuda.is_available() = False
elif device: # non-cpu device requested
+ if device == 'cuda':
+ 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(',', ''))):