Add values for testing on CUDA 12.1 image

Author: nitins17

.github/workflows/pytest_gpu.yml

@@ -16,16 +16,22 @@ on:
         - 'no'
 
 jobs:
-  build:
+  run_pytest_gpu:
     strategy:
       matrix:
-        runner: ["linux-x86-g2-48-l4-4gpu"]
+        test_env: [
+          {cuda_version: "12.3", runner: "linux-x86-g2-48-l4-4gpu", 
+          image: "gcr.io/tensorflow-testing/nosla-cuda12.3-cudnn9.1-ubuntu20.04-manylinux2014-multipython:latest"},
+          {cuda_version: "12.1", runner: "linux-x86-g2-48-l4-4gpu", 
+          image: "gcr.io/tensorflow-testing/nosla-cuda12.1-cudnn9.1-ubuntu20.04-manylinux2014-multipython:latest"},
+          ]
         python: ["3.10"]
 
-    runs-on: ${{ matrix.runner }}
+    runs-on: ${{ matrix.test_env.runner }}
     container:
-      image: "gcr.io/tensorflow-testing/nosla-cuda12.3-cudnn9.1-ubuntu20.04-manylinux2014-multipython:latest"
-
+      image: ${{ matrix.test_env.image }}
+    
+    name: "Pytest GPU (Build wheels on CUDA 12.3 and test on CUDA ${{ matrix.test_env.cuda_version }})"
     env:
       JAXCI_CLONE_MAIN_XLA: 1
       JAXCI_HERMETIC_PYTHON_VERSION: ${{ matrix.python }}