From 9f9acb2ed14ab7ac28f464245e891f856508b0fe Mon Sep 17 00:00:00 2001 From: Somdev Sangwan Date: Thu, 1 Jun 2023 15:07:44 +0530 Subject: [PATCH] dont break other gpus --- run.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/run.py b/run.py index 6a75097..abca34b 100755 --- a/run.py +++ b/run.py @@ -70,7 +70,8 @@ def pre_check(): if not os.path.isfile(model_path): quit('File "inswapper_128.onnx" does not exist!') if '--gpu' in sys.argv: - if 'ROCMExecutionProvider' not in core.globals.providers: + NVIDIA_PROVIDERS = ['CUDAExecutionProvider', 'TensorrtExecutionProvider'] + if len(list(set(core.globals.providers) - set(NVIDIA_PROVIDERS))) == 1: CUDA_VERSION = torch.version.cuda CUDNN_VERSION = torch.backends.cudnn.version() if not torch.cuda.is_available() or not CUDA_VERSION: