Update base.py (#39)

added ROCM
Tau-J · Nov 18, 2024 · 875c773 · 875c773
1 parent b9802de
commit 875c773
Showing 1 changed file with 1 addition and 1 deletion.
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,