diff --git a/tools/ci_build/github/azure-pipelines/orttraining-pai-ci-pipeline.yml b/tools/ci_build/github/azure-pipelines/orttraining-pai-ci-pipeline.yml
index a13e380af6..42e5fc5269 100644
--- a/tools/ci_build/github/azure-pipelines/orttraining-pai-ci-pipeline.yml
+++ b/tools/ci_build/github/azure-pipelines/orttraining-pai-ci-pipeline.yml
@@ -84,7 +84,7 @@ jobs:
     inputs:
       script: |-
         export KERNEL_EXPLORER_BUILD_DIR=./build/Release
-        pytest ./onnxruntime/python/tools/kernel_explorer/
+        pytest ./onnxruntime/python/tools/kernel_explorer/ -n 16 --reruns 1
     displayName: 'Run kernel explorer tests'
     condition: and(succeededOrFailed(), eq(variables.onnxruntimeBuildSucceeded, 'true'))
 
diff --git a/tools/ci_build/github/pai/rocm-ci-pipeline-env.Dockerfile b/tools/ci_build/github/pai/rocm-ci-pipeline-env.Dockerfile
index 106573a30b..bdee3cc0cd 100644
--- a/tools/ci_build/github/pai/rocm-ci-pipeline-env.Dockerfile
+++ b/tools/ci_build/github/pai/rocm-ci-pipeline-env.Dockerfile
@@ -37,7 +37,9 @@ RUN pip install \
       sentencepiece \
       dill==0.3.4 \
       wget \
-      pytorch_lightning==1.6.0
+      pytorch_lightning==1.6.0 \
+      pytest-xdist \
+      pytest-rerunfailures
 
 RUN pip install torch-ort --no-dependencies
 ENV ORTMODULE_ONNX_OPSET_VERSION=15