try with faster docker

ghstack-source-id: 5138a0fb309b05943e0746f911d7bdcb77f016e5 Pull Request resolved: #329
pytorch · May 15, 2024 · 9ea0571 · 9ea0571
1 parent 96a5cdf
commit 9ea0571
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/.github/workflows/unit_test_4gpu.yaml b/.github/workflows/unit_test_4gpu.yaml
@@ -17,7 +17,7 @@ jobs:
       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
-      # docker-image: "pytorch/pytorch:2.3.0-cuda12.1-cudnn8-runtime"
+      docker-image: "pytorch/pytorch:2.3.0-cuda12.1-cudnn8-runtime"
       repository: "pytorch/torchtitan"
       upload-artifact: "outputs"
       script: |