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

.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