try this new flag

keras-team · Feb 5, 2025 · 96d3677 · 96d3677
1 parent 74f569e
commit 96d3677
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/.kokoro/github/ubuntu/gpu/build.sh b/.kokoro/github/ubuntu/gpu/build.sh
@@ -43,7 +43,7 @@ fi
 if [ "$KERAS_BACKEND" == "jax" ]
 then
    echo "JAX backend detected."
-   export XLA_PYTHON_CLIENT_PREALLOCATE=false
+   export XLA_FLAGS=--xla_gpu_graph_level=0
    pip install -r requirements-jax-cuda.txt --progress-bar off --timeout 1000
    pip uninstall -y keras keras-nightly
    python3 -c 'import jax;print(jax.__version__);print(jax.default_backend())'