diff --git a/.github/actions/setup-venv/action.yml b/.github/actions/setup-venv/action.yml
index 1006ca91f..7179143dd 100644
--- a/.github/actions/setup-venv/action.yml
+++ b/.github/actions/setup-venv/action.yml
@@ -45,7 +45,6 @@ runs:
         pip install 'torch${{ inputs.torch-version }}' --extra-index-url https://download.pytorch.org/whl/cpu
         pip install -e .[all]
         pip install -e hf_olmo
-        pip install --no-cache-dir triton==2.0.0 https://storage.googleapis.com/ai2-python-wheels/flash_attn/flash_attn-0.2.8%2Bcu118torch2.0.0-cp310-cp310-linux_x86_64.whl
 
 
     - if: steps.virtualenv-cache.outputs.cache-hit == 'true'
@@ -55,7 +54,6 @@ runs:
         . .venv/bin/activate
         pip install --no-deps -e .[all]
         pip install --no-deps -e hf_olmo
-        pip install --no-deps --no-cache-dir triton==2.0.0 https://storage.googleapis.com/ai2-python-wheels/flash_attn/flash_attn-0.2.8%2Bcu118torch2.0.0-cp310-cp310-linux_x86_64.whl
 
     - shell: bash
       run: |
diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml
index 7fd0d7b91..ab1cf6523 100644
--- a/.github/workflows/main.yml
+++ b/.github/workflows/main.yml
@@ -128,6 +128,10 @@ jobs:
         run: |
           echo "COMMIT_SHA=$GITHUB_SHA" >> $GITHUB_ENV
 
+      - name: Install flash attn
+        - run: |
+            pip install --no-cache-dir triton==2.0.0 https://storage.googleapis.com/ai2-python-wheels/flash_attn/flash_attn-0.2.8%2Bcu118torch2.0.0-cp310-cp310-linux_x86_64.whl
+
       - name: GPU Tests
         uses: allenai/beaker-run-action@v1.2
         if: env.BEAKER_TOKEN != ''