Update base.py (#39)

jschoch · web-flow · commit 875c773bce0a · 2024-11-18T09:46:59.000+08:00
added ROCM
diff --git a/rtmlib/tools/base.py b/rtmlib/tools/base.py
@@ -24,11 +24,11 @@ def check_mps_support():
     'onnxruntime': {
         'cpu': 'CPUExecutionProvider',
         'cuda': 'CUDAExecutionProvider',
+        'rocm': 'ROCMExecutionProvider',
         'mps': 'CoreMLExecutionProvider' if check_mps_support() else 'CPUExecutionProvider'
     },
 }
 
-
 class BaseTool(metaclass=ABCMeta):
 
     def __init__(self,
