From b12db158d22b280cd754eee4b0a7655e86da71a5 Mon Sep 17 00:00:00 2001
From: frostedoyster <bigi.f@libero.it>
Date: Wed, 14 Feb 2024 14:03:24 +0100
Subject: [PATCH] More torch versions

---
 .github/workflows/build-cuda-wheels.yml   |  8 ++++----
 sphericart-torch/build-backend/backend.py | 11 +++--------
 2 files changed, 7 insertions(+), 12 deletions(-)

diff --git a/.github/workflows/build-cuda-wheels.yml b/.github/workflows/build-cuda-wheels.yml
index 2587ff056..23afe7f82 100644
--- a/.github/workflows/build-cuda-wheels.yml
+++ b/.github/workflows/build-cuda-wheels.yml
@@ -19,12 +19,12 @@ jobs:
     strategy:
       matrix:
         os: [ubuntu-20.04]
-        torch: ["2.2.0"]
+        torch: ["1.13.0", "2.0.0", "2.1.0", "2.2.0"]
         include:
           - cuda: "11.8"
-            torch_cuda: "118"
+            torch_index_cuda: "118"
           - cuda: "12.1"
-            torch_cuda: "121"
+            torch_index_cuda: "121"
     defaults:
       run:
         shell: pwsh
@@ -73,7 +73,7 @@ jobs:
           CIBW_SKIP: "*-musllinux* *-win32 *-manylinux_i686"
           CIBW_MANYLINUX_X86_64_IMAGE: manylinux2014
           # set environment variables for sphericart-torch build
-          CIBW_ENVIRONMENT: SPHERICART_ARCH_NATIVE=OFF CUDACXX=/usr/local/cuda/bin/nvcc TORCH_CUDA_ARCH_LIST=All CUDAARCHS=all SPHERICART_TORCH_TORCH_VERSION=${{ matrix.torch }} PIP_EXTRA_INDEX_URL=https://download.pytorch.org/whl/cu${{ matrix.torch_cuda }}
+          CIBW_ENVIRONMENT: SPHERICART_ARCH_NATIVE=OFF CUDACXX=/usr/local/cuda/bin/nvcc TORCH_CUDA_ARCH_LIST=All CUDAARCHS=all SPHERICART_TORCH_TORCH_VERSION=${{ matrix.torch }} PIP_EXTRA_INDEX_URL=https://download.pytorch.org/whl/cu${{ matrix.torch_index_cuda }}
           # do not complain for missing libtorch.so in sphericart-torch wheel
           CIBW_REPAIR_WHEEL_COMMAND_LINUX: |
             auditwheel repair --exclude libtorch.so --exclude libtorch_cpu.so --exclude libtorch_cuda.so --exclude libc10.so --exclude libc10_cuda.so -w {dest_dir} {wheel}
diff --git a/sphericart-torch/build-backend/backend.py b/sphericart-torch/build-backend/backend.py
index e6f03b5a0..2d60b9697 100644
--- a/sphericart-torch/build-backend/backend.py
+++ b/sphericart-torch/build-backend/backend.py
@@ -1,19 +1,14 @@
-# this is a custom Python build backend wrapping setuptool's to add a build-time
-# dependencies to metatensor-core, using the local version if it exists, and
-# otherwise falling back to the one on PyPI.
+# this is a custom Python build backend wrapping setuptool's to set a
+# specific torch version as a build dependency, based on an environment
+# variable
 import os
 
 from setuptools import build_meta
 
 TORCH_VERSION = os.environ.get("SPHERICART_TORCH_TORCH_VERSION")
-# CUDA_VERSION = os.environ.get("SPHERICART_TORCH_CUDA_VERSION")
 
 if TORCH_VERSION is not None:
-    # force a specific version of torch+cuda
     TORCH_DEP = f"torch =={TORCH_VERSION}"
-    # if CUDA_VERSION is not None:
-    #     extra_index_url = f" --index-url https://download.pytorch.org/whl/cu{CUDA_VERSION.replace('.', '')}"
-    #     TORCH_DEP += extra_index_url
 else:
     TORCH_DEP = "torch >=1.13"