Only default to cpu on auto or current.

Not when a specific device was requested.
This commit is contained in:
Marc-Andre Ferland 2022-10-28 01:08:00 -04:00
parent 0a569146a8
commit 71c6beadb4

View File

@ -112,9 +112,12 @@ def device_init(device_selection=None):
thread_data.device = 'cpu' thread_data.device = 'cpu'
return return
if not torch.cuda.is_available(): if not torch.cuda.is_available():
print('WARNING: torch.cuda is not available. Using the CPU, but this will be very slow!') if device_selection == 'auto' or device_selection == 'current':
thread_data.device = 'cpu' print('WARNING: torch.cuda is not available. Using the CPU, but this will be very slow!')
return thread_data.device = 'cpu'
return
else:
raise EnvironmentError('torch.cuda is not available.')
device_count = torch.cuda.device_count() device_count = torch.cuda.device_count()
if device_count <= 1 and device_selection == 'auto': if device_count <= 1 and device_selection == 'auto':
device_selection = 'current' # Use 'auto' only when there is more than one compatible device found. device_selection = 'current' # Use 'auto' only when there is more than one compatible device found.