use released image name

mosaicml · Sep 10, 2024 · 51c9240 · 51c9240
1 parent 4cdaf9c
commit 51c9240
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 5 deletions.
diff --git a/.github/workflows/pr-cpu.yaml b/.github/workflows/pr-cpu.yaml
@@ -25,11 +25,11 @@ jobs:
           markers: not daily and not remote and not gpu and not doctest
           pytest_command: coverage run -m pytest
         - name: cpu-3.11-2.4
-          container: ghcr.io/databricks-mosaic/ci-staging:c8adabd3-6859-44d6-8bc6-6d6f2c69000e 
+          container: mosaicml/pytorch:2.4.1_cpu-python3.11-ubuntu20.04 
           markers: not daily and not remote and not gpu and not doctest
           pytest_command: coverage run -m pytest
         - name: cpu-doctest
-          container: mosaicml/pytorch:2.4.0_cpu-python3.11-ubuntu20.04
+          container: mosaicml/pytorch:2.4.1_cpu-python3.11-ubuntu20.04
           markers: not daily and not remote and not gpu and doctest
           pytest_command: coverage run -m pytest tests/test_docs.py
     steps:

diff --git a/.github/workflows/pr-gpu.yaml b/.github/workflows/pr-gpu.yaml
@@ -16,7 +16,7 @@ jobs:
       matrix:
         include:
         - name: gpu-3.11-2.4-1
-          container: ghcr.io/databricks-mosaic/ci-staging:0ccca195-2605-483b-93a1-4753b5d53669 
+          container: mosaicml/pytorch:2.4.1_cu124-python3.11-ubuntu20.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.11-2.4-2
-          container: ghcr.io/databricks-mosaic/ci-staging:0ccca195-2605-483b-93a1-4753b5d53669
+          container: mosaicml/pytorch:2.4.1_cu124-python3.11-ubuntu20.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.11-2.4-4
-          container: ghcr.io/databricks-mosaic/ci-staging:0ccca195-2605-483b-93a1-4753b5d53669 
+          container: mosaicml/pytorch:2.4.1_cu124-python3.11-ubuntu20.04 
           markers: not daily and not remote and gpu and (doctest or not doctest)
           pytest_command: coverage run -m pytest
           composer_package_name: mosaicml