run the bazel gpu rbe test on linux cpu runners

Author: nitins17

.github/workflows/bazel_gpu_rbe.yml

@@ -19,11 +19,11 @@ jobs:
   build:
     strategy:
       matrix:
-        runner: ["linux-x86-g2-16-l4-1gpu"]
+        runner: ["linux-x86-n2-16"]
 
     runs-on: ${{ matrix.runner }}
     container:
-      image: "gcr.io/tensorflow-testing/nosla-cuda12.3-cudnn9.1-ubuntu20.04-manylinux2014-multipython:latest"
+      image: "us-central1-docker.pkg.dev/tensorflow-sigs/tensorflow/ml-build:latest"
 
     env:
       JAXCI_CLONE_MAIN_XLA: 1

ci/run_bazel_test_gpu_rbe.sh

@@ -34,12 +34,10 @@ fi
 # Set up the build environment.
 source "ci/utilities/setup_build_environment.sh"
 
-# Run Bazel GPU tests with RBE (single accelerator tests)
-nvidia-smi
+# Run Bazel GPU tests with RBE (single accelerator tests with one GPU apiece.)
 echo "Running RBE GPU tests..."
 
 # Only Linux x86 builds run GPU tests
-# Runs single accelerator tests with one GPU apiece.
 bazel test --config=rbe_linux_x86_64_cuda \
     --repo_env=HERMETIC_PYTHON_VERSION="$JAXCI_HERMETIC_PYTHON_VERSION" \
     --override_repository=xla="${JAXCI_XLA_GIT_DIR}" \