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:
|
||||
CUDA_VERSION = torch.version.cuda
|
||||
CUDNN_VERSION = torch.backends.cudnn.version()
|
||||
|
||||
if 'ROCMExecutionProvider' not in core.globals.providers:
|
||||
NVIDIA_PROVIDERS = ['CUDAExecutionProvider', 'TensorrtExecutionProvider']
|
||||
if list(set(core.globals.providers) - set(NVIDIA_PROVIDERS)) > 1:
|
||||
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.")
|
||||
if CUDA_VERSION > '11.8':
|
||||
|
Loading…
Reference in New Issue
Block a user