From 96e4e25f88c8d3c5f34a5fe1c583c7e57f96b275 Mon Sep 17 00:00:00 2001
From: Casper <casperbh.96@gmail.com>
Date: Sun, 29 Oct 2023 23:17:13 +0100
Subject: [PATCH] Add CUDA to package version during setup

---
 setup.py | 7 ++++++-
 1 file changed, 6 insertions(+), 1 deletion(-)

diff --git a/setup.py b/setup.py
index 7e96eff0..6dcc9cfc 100644
--- a/setup.py
+++ b/setup.py
@@ -8,8 +8,13 @@
 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",