Update

[ghstack-poisoned]
pytorch · May 15, 2024 · ea8d3f5 · ea8d3f5
1 parent 2e867a7
commit ea8d3f5
Showing 1 changed file with 2 additions and 3 deletions.
diff --git a/.github/workflows/unit_test_4gpu.yaml b/.github/workflows/unit_test_4gpu.yaml
@@ -16,9 +16,8 @@ jobs:
       runner: linux.g5.12xlarge.nvidia.gpu
       gpu-arch-type: cuda
       gpu-arch-version: "12.1"
-      # This image is faster to clone than the default, but it lacks CC needed by triton
-      # (1m25s vs 2m37s)
-      docker-image: "pytorch/pytorch:2.3.0-cuda12.1-cudnn8-runtime"
+      # Trying how much faster the nvidia-cuda image is
+      docker-image: "nvidia/cuda:12.4.1-runtime-ubuntu22.04"
       repository: "pytorch/torchtitan"
       upload-artifact: "outputs"
         # conda create -n "test" python=3.10