diff --git a/.github/workflows/wheel_tests_continuous.yml b/.github/workflows/wheel_tests_continuous.yml
index 207075fd0340..d5de5a01103a 100644
--- a/.github/workflows/wheel_tests_continuous.yml
+++ b/.github/workflows/wheel_tests_continuous.yml
@@ -22,6 +22,10 @@ permissions:
       contents: read
 
 on:
+  # Add a pull request for testing
+  pull_request:
+    branches:
+      - main
   schedule:
     - cron: "0 */3 * * *" # Run once every 3 hours
   workflow_dispatch: # allows triggering the workflow run manually
@@ -145,7 +149,7 @@ jobs:
     strategy:
         fail-fast: false # don't cancel all jobs on failure
         matrix:
-          runner: ["linux-x86-n2-16", "linux-arm64-t2a-48"]
+          runner: ["linux-x86-n2-16", "linux-arm64-c4a-64"]
           python: ["3.10",]
           enable-x64: [1, 0]
     name: "Bazel CPU tests with ${{ format('{0}', 'py_import') }}"
diff --git a/ci/run_bazel_test_cpu_py_import_rbe.sh b/ci/run_bazel_test_cpu_py_import_rbe.sh
index 9a17397c47ff..2d5d4748af21 100755
--- a/ci/run_bazel_test_cpu_py_import_rbe.sh
+++ b/ci/run_bazel_test_cpu_py_import_rbe.sh
@@ -50,6 +50,7 @@ if [[ $os == "darwin" ]] || ( [[ $os == "linux" ]] && [[ $arch == "aarch64" ]] )
             --action_env=JAX_ENABLE_X64="$JAXCI_ENABLE_X64" \
             --test_output=errors \
             --color=yes \
+            --test_timeout=500 \
             --strategy=TestRunner=local \
             --//jax:build_jaxlib=wheel \
             --//jax:build_jax=wheel \