From 7514aae76afe1861fe6d6e4734cf8518c988ec3c Mon Sep 17 00:00:00 2001
From: iefgnoix <isaacwxf23@gmail.com>
Date: Fri, 17 May 2024 18:59:34 +0000
Subject: [PATCH] fix more

---
 .github/workflows/_test.yml | 7 +++++++
 1 file changed, 7 insertions(+)

diff --git a/.github/workflows/_test.yml b/.github/workflows/_test.yml
index 55f5641a514..bbcb185d203 100644
--- a/.github/workflows/_test.yml
+++ b/.github/workflows/_test.yml
@@ -92,6 +92,13 @@ jobs:
         with:
           name: torch-xla-wheels
           path: /tmp/wheels/
+        if: ! ${{ inputs.require-pytorch-cuda }}
+      - name: Fetch wheels (pytorch and pytorch/xla with CUDA enabled) if needed.
+        uses: actions/download-artifact@v4
+        with:
+          name: torch-with-cuda-xla-with-cuda-wheels
+          path: /tmp/wheels/
+        if: ${{ inputs.require-pytorch-cuda }}
       - name: Fetch CPP test binaries
         uses: actions/download-artifact@v4
         with: