feat(ml): introduce support of onnxruntime-rocm for AMD GPU

2024-07-13 00:40:29 +02:00
parent 3f4bbab4eb
commit fe26ccd1b7
14 changed files with 275 additions and 78 deletions
@@ -88,7 +88,7 @@ class OrtSession:
            match provider:
                case "CPUExecutionProvider":
                    options = {"arena_extend_strategy": "kSameAsRequested"}
-                case "CUDAExecutionProvider":
+                case "CUDAExecutionProvider" | "ROCMExecutionProvider":
                    options = {"arena_extend_strategy": "kSameAsRequested", "device_id": settings.device_id}
                case "OpenVINOExecutionProvider":
                    options = {