diff --git a/.github/workflows/docker.yaml b/.github/workflows/docker.yaml index 298c1e0d80..76874f5440 100644 --- a/.github/workflows/docker.yaml +++ b/.github/workflows/docker.yaml @@ -27,7 +27,7 @@ jobs: base_image: mosaicml/pytorch:2.1.0_cu121-python3.10-ubuntu20.04 dep_groups: '[gpu]' - name: '2.1.0_cu121_aws' - base_image: mosaicml/gs + base_image: mosaicml/pytorch:2.1.0_cu121-python3.10-ubuntu20.04-aws dep_groups: '[gpu]' - name: '2.1.0_cu121_flash2' base_image: mosaicml/pytorch:2.1.0_cu121-python3.10-ubuntu20.04