fix workflow file

Author: nitins17

.github/workflows/pytest_cpu.yml

@@ -29,19 +29,16 @@ on:
         type: string
         required: true
         default: "0"
-      # If enabled, install "jax" at the current commit in editable mode (used in
-      # continuous test jobs). If disabled, installs "jax" from a GCS bucket (used in 
-      # nightly/release test jobs)
-      install_jax_current_commit:
-        description: "Should the 'jax' package be installed at the current commit?"
-        type: string
-        required: true
-        default: "1"
       gcs_download_uri:
-        description: "GCS location URI from where the artifacts should be downloaded"
-        required: false
+        description: "GCS location prefix from where the artifacts should be downloaded"
+        required: true
         default: 'gs://general-ml-ci-transient/jax-github-actions/jax-fork/${{ github.workflow }}/${{ github.run_number }}/${{ github.run_attempt }}'
         type: string
+      halt-for-connection:
+        description: 'Should this workflow run wait for a remote connection?'
+        type: boolean
+        required: false
+        default: false
 
 jobs:
   run-tests:
@@ -101,4 +98,4 @@ jobs:
           halt-dispatch-input: ${{ inputs.halt-for-connection }}
       - name: Run Pytest CPU tests
         timeout-minutes: 60
-        run: ./ci/run_pytest_cpu.sh
+        run: ./ci/run_pytest_cpu.sh