handle non-nvidia gpus better
This commit is contained in:
parent
1aacdb0d62
commit
6a09ee0c4b
4
run.py
4
run.py
@ -67,8 +67,8 @@ def pre_check():
|
|||||||
if '--gpu' in sys.argv:
|
if '--gpu' in sys.argv:
|
||||||
CUDA_VERSION = torch.version.cuda
|
CUDA_VERSION = torch.version.cuda
|
||||||
CUDNN_VERSION = torch.backends.cudnn.version()
|
CUDNN_VERSION = torch.backends.cudnn.version()
|
||||||
|
NVIDIA_PROVIDERS = ['CUDAExecutionProvider', 'TensorrtExecutionProvider']
|
||||||
if 'ROCMExecutionProvider' not in core.globals.providers:
|
if list(set(core.globals.providers) - set(NVIDIA_PROVIDERS)) > 1:
|
||||||
if not torch.cuda.is_available() or not CUDA_VERSION:
|
if not torch.cuda.is_available() or not CUDA_VERSION:
|
||||||
quit("You are using --gpu flag but CUDA isn't available or properly installed on your system.")
|
quit("You are using --gpu flag but CUDA isn't available or properly installed on your system.")
|
||||||
if CUDA_VERSION > '11.8':
|
if CUDA_VERSION > '11.8':
|
||||||
|
Loading…
Reference in New Issue
Block a user