Update ROCm to not hang when running on complex RNN models

This commit is contained in:
Nicolas Mowen 2025-09-17 17:33:26 -06:00
parent 1efff67e32
commit 33042592c7
2 changed files with 10 additions and 0 deletions

View File

@ -67,6 +67,9 @@ FROM deps-prelim AS rocm-prelim-hsa-override0
ENV HSA_ENABLE_SDMA=0
ENV TF_ROCM_USE_IMMEDIATE_MODE=1
# avoid kernel crashes
ENV HIP_FORCE_DEV_KERNARG=1
COPY --from=rocm-dist / /
RUN ldconfig

View File

@ -51,6 +51,13 @@ class ONNXDetector(DetectionApi):
"enable_cuda_graph": True,
}
sess_options = None
if providers[0] == "ROCMExecutionProvider":
# avoid AMD GPU kernel crashes
sess_options = ort.SessionOptions()
sess_options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_DISABLE_ALL
self.model = ort.InferenceSession(
path, providers=providers, provider_options=options
)