Using cu128 instead of cu126 for pr-gpu and daily tests (#3867)

Author: rithwik-db

.github/workflows/daily.yaml

@@ -104,7 +104,7 @@ jobs:
           composer_package_name: "mosaicml"
           gpu_num: 1
         - name: "gpu-3.12-2.7-1-gpu"
-          container: mosaicml/pytorch:2.7.0_cu126-python3.12-ubuntu22.04
+          container: mosaicml/pytorch:2.7.0_cu128-python3.12-ubuntu22.04
           markers: "(daily or not daily) and (remote or not remote) and gpu and (doctest or not doctest)"
           pytest_command: "coverage run -m pytest"
           composer_package_name: "mosaicml"
@@ -116,7 +116,7 @@ jobs:
           composer_package_name: "mosaicml"
           gpu_num: 2
         - name: "gpu-3.12-2.7-2-gpu"
-          container: mosaicml/pytorch:2.7.0_cu126-python3.12-ubuntu22.04
+          container: mosaicml/pytorch:2.7.0_cu128-python3.12-ubuntu22.04
           markers: "(daily or not daily) and (remote or not remote) and gpu and (doctest or not doctest)"
           pytest_command: "coverage run -m pytest"
           composer_package_name: "mosaicml"
@@ -128,7 +128,7 @@ jobs:
           composer_package_name: "mosaicml"
           gpu_num: 4
         - name: "gpu-3.12-2.7-4-gpu"
-          container: mosaicml/pytorch:2.7.0_cu126-python3.12-ubuntu22.04
+          container: mosaicml/pytorch:2.7.0_cu128-python3.12-ubuntu22.04
           markers: "(daily or not daily) and (remote or not remote) and gpu and (doctest or not doctest)"
           pytest_command: "coverage run -m pytest"
           composer_package_name: "mosaicml"

.github/workflows/pr-gpu.yaml

@@ -16,7 +16,7 @@ jobs:
       matrix:
         include:
         - name: gpu-3.12-2.7-1
-          container: mosaicml/pytorch:2.7.0_cu126-python3.12-ubuntu22.04
+          container: mosaicml/pytorch:2.7.0_cu128-python3.12-ubuntu22.04
           markers: not daily and not remote and gpu and (doctest or not doctest)
           pytest_command: coverage run -m pytest
           composer_package_name: mosaicml
@@ -45,7 +45,7 @@ jobs:
       matrix:
         include:
         - name: gpu-3.12-2.7-2
-          container: mosaicml/pytorch:2.7.0_cu126-python3.12-ubuntu22.04
+          container: mosaicml/pytorch:2.7.0_cu128-python3.12-ubuntu22.04
           markers: not daily and not remote and gpu and (doctest or not doctest)
           pytest_command: coverage run -m pytest
           composer_package_name: mosaicml
@@ -75,7 +75,7 @@ jobs:
       matrix:
         include:
         - name: gpu-3.12-2.7-4
-          container: mosaicml/pytorch:2.7.0_cu126-python3.12-ubuntu22.04
+          container: mosaicml/pytorch:2.7.0_cu128-python3.12-ubuntu22.04
           markers: not daily and not remote and gpu and (doctest or not doctest)
           pytest_command: coverage run -m pytest
           composer_package_name: mosaicml