Revert cuda changes done to test wheel breakage

Author: nitins17

.github/workflows/pytest_cuda.yml

@@ -58,8 +58,7 @@ jobs:
     runs-on: ${{ inputs.runner }}
     # Test the oldest and newest supported CUDA versions.
     container:  ${{ (contains(inputs.cuda, '12.1') && 'us-central1-docker.pkg.dev/tensorflow-sigs/tensorflow/ml-build-cuda12.1-cudnn9.8:latest') ||
-                (contains(inputs.cuda, '12.8') && 'us-central1-docker.pkg.dev/tensorflow-sigs/tensorflow/ml-build-cuda12.8-cudnn9.8:latest') ||
-                (contains(inputs.cuda, 'CUDA pip') && 'us-central1-docker.pkg.dev/tensorflow-sigs/tensorflow/ml-build:latest') }}
+                (contains(inputs.cuda, '12.8') && 'us-central1-docker.pkg.dev/tensorflow-sigs/tensorflow/ml-build-cuda12.8-cudnn9.8:latest')}}
     name: "Pytest CUDA (${{ inputs.runner }}, CUDA ${{ inputs.cuda }}, Python ${{ inputs.python }}, x64=${{ inputs.enable-x64 }})"
     env:
       JAXCI_HERMETIC_PYTHON_VERSION: "${{ inputs.python }}"

.github/workflows/wheel_tests_continuous.yml

@@ -117,7 +117,7 @@ jobs:
           # Comment out H100 and B200 on jax-fork to save on resources
           runner: ["linux-x86-g2-48-l4-4gpu"] #, "linux-x86-a3-8g-h100-8gpu", "linux-x86-a4-224-b200-1gpu"]
           python: ["3.10",]
-          cuda: ["CUDA pip",] #"12.1", "12.8"]
+          cuda: ["12.1", "12.8"]
           enable-x64: [1, 0]
           # exclude:
           #   # H100 runs only a single config, CUDA 12.8 Enable x64 1