From ce92b406e6af8a08d4a311219c2f6325c8943bea Mon Sep 17 00:00:00 2001
From: Ashleigh Watson <145302363+OutofAi@users.noreply.github.com>
Date: Sat, 11 May 2024 03:01:25 +0100
Subject: [PATCH] Update wheels-generator.yml

---
 .github/workflows/wheels-generator.yml | 9 +--------
 1 file changed, 1 insertion(+), 8 deletions(-)

diff --git a/.github/workflows/wheels-generator.yml b/.github/workflows/wheels-generator.yml
index 4c76268..8e7459c 100644
--- a/.github/workflows/wheels-generator.yml
+++ b/.github/workflows/wheels-generator.yml
@@ -208,6 +208,7 @@ jobs:
           source $HOME/miniconda/etc/profile.d/conda.sh
           conda create -n myenv python=3.10 -y
           conda activate myenv
+          conda install pytorch==2.1.2 torchvision==0.16.2 torchaudio==2.1.2 pytorch-cuda=12.1 -c pytorch -c nvidia
         shell: bash
           
       - uses: actions/checkout@v3
@@ -227,14 +228,6 @@ jobs:
           pip install setuptools wheel
         shell: bash
 
-      - name: Install Pytorch
-        run: |
-          source $HOME/miniconda/etc/profile.d/conda.sh
-          conda activate myenv
-          python --version
-          pip install torch==2.1.2 --index-url https://download.pytorch.org/whl/cu121
-        shell: bash
-
       - name: Generate Wheels
         run: |
           source $HOME/miniconda/etc/profile.d/conda.sh