Address comments

NVIDIA · May 20, 2024 · f4cc57b · f4cc57b
1 parent 5f89e9e
commit f4cc57b
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 3 deletions.
diff --git a/.github/container/test-jax.sh b/.github/container/test-jax.sh
@@ -142,7 +142,6 @@ case "${BATTERY}" in
         JOBS=$((NGPUS * JOBS_PER_GPU))
         EXTRA_FLAGS="--local_test_jobs=${JOBS} --test_env=JAX_TESTS_PER_ACCELERATOR=${JOBS_PER_GPU} --test_env=JAX_EXCLUDE_TEST_TARGETS=PmapTest.testSizeOverflow"
         BAZEL_TARGET="${BAZEL_TARGET} //tests:gpu_tests"
-        EXTRA_CONFIG_FLAG="--enable_cuda"
         ;;
     backend-independent)
         JOBS_PER_GPU=4
@@ -185,6 +184,6 @@ pip install matplotlib
 ## Run tests
 
 cd `jax_source_dir`
-python build/build.py --configure_only ${EXTRA_CONFIG_FLAG}
+python build/build.py --configure_only --enable_cuda
 BAZEL=$(find -type f -executable -name "bazel-*")
 $BAZEL test ${BAZEL_TARGET} ${COMMON_FLAGS} ${EXTRA_FLAGS}
diff --git a/.github/workflows/_ci.yaml b/.github/workflows/_ci.yaml
@@ -296,7 +296,7 @@ jobs:
         docker run -i --shm-size=1g --gpus all \
         ${{ needs.build-jax.outputs.DOCKER_TAG_FINAL }} \
         bash <<"EOF" |& tee test-backend-independent.log
-          test-jax.sh -b backend-independent
+          test-jax.sh --build-jaxlib -b backend-independent
         EOF
         docker run -i --shm-size=1g --gpus all \
         ${{ needs.build-jax.outputs.DOCKER_TAG_FINAL }} \