From 6a09ee0c4b63f4808666c0acba50ed71058fbedd Mon Sep 17 00:00:00 2001
From: Somdev Sangwan <s0md3v@gmail.com>
Date: Wed, 31 May 2023 10:36:18 +0530
Subject: [PATCH] handle non-nvidia gpus better

---
 run.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/run.py b/run.py
index 19a10be..a9570fc 100644
--- a/run.py
+++ b/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':