mirror of
https://github.com/w-okada/voice-changer.git
synced 2025-01-23 21:45:00 +03:00
strict cuda check
This commit is contained in:
parent
b3b8a347e1
commit
d6731f4adf
@ -30,7 +30,8 @@ class DeviceManager(object):
|
|||||||
|
|
||||||
def getOnnxExecutionProvider(self, gpu: int):
|
def getOnnxExecutionProvider(self, gpu: int):
|
||||||
availableProviders = onnxruntime.get_available_providers()
|
availableProviders = onnxruntime.get_available_providers()
|
||||||
if gpu >= 0 and "CUDAExecutionProvider" in availableProviders:
|
devNum = torch.cuda.device_count()
|
||||||
|
if gpu >= 0 and "CUDAExecutionProvider" in availableProviders and devNum > 0:
|
||||||
return ["CUDAExecutionProvider"], [{"device_id": gpu}]
|
return ["CUDAExecutionProvider"], [{"device_id": gpu}]
|
||||||
elif gpu >= 0 and "DmlExecutionProvider" in availableProviders:
|
elif gpu >= 0 and "DmlExecutionProvider" in availableProviders:
|
||||||
return ["DmlExecutionProvider"], [{}]
|
return ["DmlExecutionProvider"], [{}]
|
||||||
|
Loading…
Reference in New Issue
Block a user