CUDA 11.8.0 and 12.1.1 build (#128)

.github/workflows/build.yaml

@@ -41,7 +41,7 @@ jobs:
       matrix:
         os: [ubuntu-20.04, windows-latest]
         pyver: ["3.8", "3.9", "3.10", "3.11"]
-        cuda: ["11.8"]
+        cuda: ["11.8.0", "12.1.1"]
     defaults:
       run:
         shell: pwsh
@@ -81,15 +81,17 @@ jobs:
       - name: Install Dependencies
         run: |
           # Install CUDA toolkit
-          mamba install -y 'cuda' -c "nvidia/label/cuda-${env:CUDA_VERSION}.0"
+          mamba install -y 'cuda' -c "nvidia/label/cuda-${env:CUDA_VERSION}"
 
           # Env variables
           $env:CUDA_PATH = $env:CONDA_PREFIX
           $env:CUDA_HOME = $env:CONDA_PREFIX
           
           # Install torch
-          $cudaVersionPytorch = $env:CUDA_VERSION.Replace('.', '')
-          python -m pip install --upgrade --no-cache-dir torch==2.0.1+cu$cudaVersionPytorch --index-url https://download.pytorch.org/whl/cu$cudaVersionPytorch
+          $cudaVersion = $env:CUDA_VERSION.Replace('.', '')
+          $cudaVersionPytorch = $cudaVersion.Substring(0, $cudaVersion.Length - 1)
+          if ([int]$cudaVersionPytorch -gt 118) { $pytorchVersion = "torch==2.1.0" } else {$pytorchVersion = "torch==2.0.1"}
+          python -m pip install --upgrade --no-cache-dir $pytorchVersion+cu$cudaVersionPytorch --index-url https://download.pytorch.org/whl/cu$cudaVersionPytorch
           python -m pip install build setuptools wheel ninja
 
           # Print version information

setup.py

@@ -8,8 +8,13 @@ from torch.utils.cpp_extension import BuildExtension, CUDA_HOME, CUDAExtension
 os.environ["CC"] = "g++"
 os.environ["CXX"] = "g++"
 
+try:
+    CUDA_VERSION = "".join(os.environ.get("CUDA_VERSION", torch.version.cuda).split("."))
+except Exception as ex:
+    raise RuntimeError("Your system must have an Nvidia GPU for installing AutoAWQ")
+
 common_setup_kwargs = {
-    "version": "0.1.5",
+    "version": f"0.1.5+cu{CUDA_VERSION}",
     "name": "autoawq",
     "author": "Casper Hansen",
     "license": "MIT",