diff --git a/.github/workflows/pr-gpu.yaml b/.github/workflows/pr-gpu.yaml
index 335d049306..07e811e244 100644
--- a/.github/workflows/pr-gpu.yaml
+++ b/.github/workflows/pr-gpu.yaml
@@ -25,7 +25,7 @@ jobs:
           pytest_command: "coverage run -m pytest"
           pip_deps: "[all]"
         - name: "gpu-2.3.1"
-          container: mosaicml/llm-foundry:2.3.1_cu121_flash2-latest
+          container: mosaicml/llm-foundry:2.3.1_cu121-latest
           markers: "gpu"
           pytest_command: "coverage run -m pytest"
           pip_deps: "[all]"