diff --git a/.github/workflows/Public_Jax_CI.yml b/.github/workflows/Public_Jax_CI.yml
index da13ea7..404e1fe 100644
--- a/.github/workflows/Public_Jax_CI.yml
+++ b/.github/workflows/Public_Jax_CI.yml
@@ -69,6 +69,7 @@ jobs:
         basekit_path=/home/sdp
         source ${basekit_path}/intel/oneapi/compiler/latest/env/vars.sh
         source ${basekit_path}/intel/oneapi/mkl/latest/env/vars.sh
+        export ZE_AFFINITY_MASK=0,1
         cd ~/actions-runner/workspace_openxla/intel-extension-for-openxla/intel-extension-for-openxla/jax
         commit_hash=$(git rev-parse --short HEAD)
         mkdir -p ./logs/ut