Improved detection of missing cuda:0 and added warning to console about how to fix.

This commit is contained in:
Marc-Andre Ferland
2022-10-17 02:27:30 -04:00
parent 554b67a2f0
commit c92129ac63
3 changed files with 31 additions and 9 deletions

View File

@ -130,7 +130,7 @@ def device_init(device_selection=None):
device_selection = device_selection.lower()
if device_selection.startswith('gpu:'):
device_selection = int(device_selection[4:])
if device_selection != 'cuda' and device_selection != 'current' and device_selection != 'gpu':
if device_selection != 'current' and device_selection != 'gpu':
if device_select(device_selection):
if isinstance(device_selection, int):
print(f'Setting GPU:{device_selection} as active')