diff --git a/.buildkite/gpu_pipeline.yml b/.buildkite/gpu_pipeline.yml
index bfd5329..004895a 100644
--- a/.buildkite/gpu_pipeline.yml
+++ b/.buildkite/gpu_pipeline.yml
@@ -23,6 +23,7 @@ steps:
       echo "--- :python: Test"
 
       export CUDA_DIR=`pwd`/bazel-bin/test/llama.runfiles/pypi_nvidia_cuda_nvcc_cu12/site-packages/nvidia/cuda_nvcc
+      export XLA_FLAGS=--xla_gpu_cuda_data_dir=$CUDA_DIR
       export LD_LIBRARY_PATH="`pwd`/bazel-bin/test/llama.runfiles/pypi_nvidia_cusolver_cu12/site-packages/nvidia/cusolver:$LD_LIBRARY_PATH"
       export LD_LIBRARY_PATH="`pwd`/bazel-bin/test/llama.runfiles/pypi_nvidia_cudnn_cu12/site-packages/nvidia/cudnn/lib:$LD_LIBRARY_PATH"
       export LD_LIBRARY_PATH="`pwd`/bazel-bin/test/test.runfiles/pypi_nvidia_cublas_cu12/site-packages/nvidia/cublas/lib:$LD_LIBRARY_PATH"