Add code for saving for release on PyPi

.github/workflows/build.yaml

@@ -46,6 +46,7 @@ jobs:
       run:
         shell: pwsh
     env:
+      PYPI_CUDA_VERSION: "12.1.1"
       CUDA_VERSION: ${{ matrix.cuda }}
     
     steps:
@@ -106,6 +107,11 @@ jobs:
           $env:CUDA_PATH = $env:CONDA_PREFIX
           $env:CUDA_HOME = $env:CONDA_PREFIX
 
+          # Only add +cu118 to wheel if not releasing on PyPi
+          if ( $env:CUDA_VERSION -eq $env:PYPI_CUDA_VERSION ){
+            $env:PYPI_BUILD = 1
+          }
+
           python setup.py sdist bdist_wheel
       
       - name: Upload Assets

setup.py

@@ -8,11 +8,13 @@ from torch.utils.cpp_extension import BuildExtension, CUDA_HOME, CUDAExtension
 os.environ["CC"] = "g++"
 os.environ["CXX"] = "g++"
 AUTOAWQ_KERNELS_VERSION = "0.0.1"
+PYPI_BUILD = os.getenv("PYPI_BUILD", "0") == "1"
 
-try:
+if not PYPI_BUILD:
+    try:
         CUDA_VERSION = "".join(os.environ.get("CUDA_VERSION", torch.version.cuda).split("."))[:3]
         AUTOAWQ_KERNELS_VERSION += f"+cu{CUDA_VERSION}"
-except Exception as ex:
+    except Exception as ex:
         raise RuntimeError("Your system must have an Nvidia GPU for installing AutoAWQ")
 
 common_setup_kwargs = {