diff --git a/qa/L0_jax_distributed_unittest/test.sh b/qa/L0_jax_distributed_unittest/test.sh
index 2b8b2e9feb..8b5710d839 100644
--- a/qa/L0_jax_distributed_unittest/test.sh
+++ b/qa/L0_jax_distributed_unittest/test.sh
@@ -4,6 +4,8 @@
 
 set -xe
 
+: ${TE_PATH:=/opt/transformerengine}
+
 pip install -r $TE_PATH/examples/jax/encoder/requirements.txt
 
 # Make encoder tests to have run-to-run deterministic to have the stable CI results