Update cuda versions for CI tests (#2994)

Summary:
Pull Request resolved: #2994

Set the CUDA versions supported by Pytorch as per https://pytorch.org/get-started/locally/

Reviewed By: TroyGarden

Differential Revision: D75248197

fbshipit-source-id: 309dc312d498c7ab93a665e23aec92a00680e5c3

Author: kausv

.github/workflows/unittest_ci.yml

@@ -23,47 +23,47 @@ jobs:
           - os: linux.g5.12xlarge.nvidia.gpu
             python-version: 3.9
             python-tag: "py39"
-            cuda-tag: "cu121"
+            cuda-tag: "cu126"
           - os: linux.g5.12xlarge.nvidia.gpu
             python-version: 3.9
             python-tag: "py39"
-            cuda-tag: "cu124"
+            cuda-tag: "cu128"
           - os: linux.g5.12xlarge.nvidia.gpu
             python-version: '3.10'
             python-tag: "py310"
             cuda-tag: "cu118"
           - os: linux.g5.12xlarge.nvidia.gpu
             python-version: '3.10'
             python-tag: "py310"
-            cuda-tag: "cu121"
+            cuda-tag: "cu126"
           - os: linux.g5.12xlarge.nvidia.gpu
             python-version: '3.10'
             python-tag: "py310"
-            cuda-tag: "cu124"
+            cuda-tag: "cu128"
           - os: linux.g5.12xlarge.nvidia.gpu
             python-version: '3.11'
             python-tag: "py311"
             cuda-tag: "cu118"
           - os: linux.g5.12xlarge.nvidia.gpu
             python-version: '3.11'
             python-tag: "py311"
-            cuda-tag: "cu121"
+            cuda-tag: "cu126"
           - os: linux.g5.12xlarge.nvidia.gpu
             python-version: '3.11'
             python-tag: "py311"
-            cuda-tag: "cu124"
+            cuda-tag: "cu128"
           - os: linux.g5.12xlarge.nvidia.gpu
             python-version: '3.12'
             python-tag: "py312"
             cuda-tag: "cu118"
           - os: linux.g5.12xlarge.nvidia.gpu
             python-version: '3.12'
             python-tag: "py312"
-            cuda-tag: "cu121"
+            cuda-tag: "cu126"
           - os: linux.g5.12xlarge.nvidia.gpu
             python-version: '3.12'
             python-tag: "py312"
-            cuda-tag: "cu124"
+            cuda-tag: "cu128"
     uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main
     permissions:
       id-token: write