Chiebot-Mirror
/
ultralytics
mirror of https://github.com/ultralytics/ultralytics.git
9
0
Fork 1
Code Issues Packages Projects Releases Wiki Activity

Add `device=cuda` support (#3133)

Browse Source
pull/3136/head
Glenn Jocher 1 year ago committed by GitHub
parent ff91fbd9c3
commit 8e60fc9276
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 2 additions and 0 deletions
Whitespace
Show all changes Ignore whitespace when comparing lines Ignore changes in amount of whitespace Ignore changes in whitespace at EOL
Unified View
Diff Options
Show Stats Download Patch File Download Diff File
  1. 2
      ultralytics/yolo/utils/torch_utils.py

2
ultralytics/yolo/utils/torch_utils.py
Unescape View File

@ -64,6 +64,8 @@ def select_device(device='', batch=0, newline=False, verbose=True):
if cpu or mps: if cpu or mps:
os.environ['CUDA_VISIBLE_DEVICES'] = '-1' # force torch.cuda.is_available() = False os.environ['CUDA_VISIBLE_DEVICES'] = '-1' # force torch.cuda.is_available() = False
elif device: # non-cpu device requested elif device: # non-cpu device requested
if device == 'cuda':
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.replace(',', ''))):

Write Preview
Loading…
Cancel
Save