From 0e539afd9fa57e66acdddbc15108073f68904c2b Mon Sep 17 00:00:00 2001
From: Ziwen Liu <ziwen.liu@czbiohub.org>
Date: Fri, 23 Feb 2024 23:55:43 -0800
Subject: [PATCH] CI: only install CPU wheels for torch

---
 .github/workflows/pr.yml | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/.github/workflows/pr.yml b/.github/workflows/pr.yml
index 6c46ee0d..f312f5d5 100644
--- a/.github/workflows/pr.yml
+++ b/.github/workflows/pr.yml
@@ -17,6 +17,7 @@ jobs:
           src: "viscy"
 
   test:
+    name: Test
     needs: [lint]
     runs-on: ubuntu-latest
     strategy:
@@ -30,6 +31,7 @@ jobs:
       - name: Install dependencies
         run: |
           python -m pip install --upgrade pip
-          pip install ".[metrics,dev]"
+          # Install cpu wheels only to speed up the build
+          pip install ".[metrics,dev]" --index-url https://download.pytorch.org/whl/cpu
       - name: Test with pytest
         run: pytest -v