diff --git a/.github/workflows/build.yaml b/.github/workflows/build.yaml
index a3931eab..ee8fc120 100644
--- a/.github/workflows/build.yaml
+++ b/.github/workflows/build.yaml
@@ -126,17 +126,11 @@ jobs:
             docker tag ghcr.io/predibase/lorax:main "$tag"
           done <<< "$tags"
 
-      - name: Grab first tag
-        env:
-          tags: ${{ steps.meta.outputs.tags }}
-        run: |
-          first_tag=$(echo "$tags" | head -n 1)
-          echo "TEST_IMAGE_TAG=$first_tag" >> $GITHUB_OUTPUT
 
       - name: Run integration tests
         uses: ./.github/workflows/integration-tests
         with:
-          test_image_tag: ${{ steps.vars.outputs.test_image_tag }}
+          test_image_tag: ${{ steps.meta.outputs.tags }}
           use_local_image: true
           github_token: ${{ secrets.GHCR_PAT }}
           huggingface_token: ${{ secrets.HUGGING_FACE_HUB_TOKEN }}
diff --git a/.github/workflows/integration-tests/action.yaml b/.github/workflows/integration-tests/action.yaml
index a52d15ba..c7492d46 100644
--- a/.github/workflows/integration-tests/action.yaml
+++ b/.github/workflows/integration-tests/action.yaml
@@ -57,7 +57,8 @@ runs:
       shell: bash
       run: |
         echo "HUGGING_FACE_HUB_TOKEN=${{ inputs.huggingface_token }}" >> $GITHUB_ENV
-        echo "TEST_IMAGE_TAG=${{ inputs.test_image_tag }}" >> $GITHUB_ENV
+        first_tag=$(echo "${{ inputs.test_image_tag }}" | head -n 1)
+        echo "TEST_IMAGE_TAG=$first_tag" >> $GITHUB_ENV
 
     - name: Run Embedding tests
       shell: bash