commit change

mosaicml · Feb 13, 2024 · 4d11c82 · 4d11c82
1 parent 7f9c425
commit 4d11c82
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/docker/generate_build_matrix.py b/docker/generate_build_matrix.py
@@ -228,14 +228,14 @@ def _main():
         'AWS_OFI_NCCL_VERSION': 'v1.7.4-aws',
         'BASE_IMAGE': 'nvidia/cuda:12.1.0-cudnn8-devel-ubuntu20.04',
         'CUDA_VERSION': '12.1.0',
-        'IMAGE_NAME': 'torch-nightly-2-3-0-20240110-cu121-python3-10',
+        'IMAGE_NAME': 'torch-nightly-2-3-0-20240110-cu121-python3-10-aws',
         'MOFED_VERSION': '',
         'NVIDIA_REQUIRE_CUDA_OVERRIDE': _get_cuda_override('12.1.0'),
         'PYTHON_VERSION': '3.10',
         'PYTORCH_VERSION': '2.3.0',
         'PYTORCH_NIGHTLY_URL': 'https://download.pytorch.org/whl/nightly/cu121',
         'PYTORCH_NIGHTLY_VERSION': 'dev20240110+cu121',
-        'TAGS': ['mosaicml/pytorch:2.3.0_cu121-nightly20240110-python3.10-ubuntu20.04'],
+        'TAGS': ['mosaicml/pytorch:2.3.0_cu121-nightly20240110-python3.10-ubuntu20.04-aws'],
         'TARGET': 'pytorch_stage',
         'TORCHVISION_VERSION': '0.18.0'
     }