CPU tests image fix (#3409)

* Revert "Optionally use `flash-attn`'s CE loss for metrics (#3394)" This reverts commit 2cf9262. revert dat boi * remove * slamm
mosaicml · Jul 21, 2024 · 0d6ef26 · 0d6ef26
1 parent 0a1a6a4
commit 0d6ef26
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/.github/workflows/pr-cpu.yaml b/.github/workflows/pr-cpu.yaml
@@ -22,7 +22,7 @@ jobs:
           markers: not daily and not remote and not gpu and not doctest
           pytest_command: coverage run -m pytest
         - name: cpu-3.11-2.3
-          container: mosaicml/pytorch:2.3.1_cu121-python3.11-ubuntu20.04
+          container: mosaicml/pytorch:2.3.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