modify setup.py

NetEase-FuXi · Dec 6, 2024 · a6c5110 · a6c5110
1 parent 358cb8c
commit a6c5110
Show file tree

Hide file tree

Showing 3 changed files with 3 additions and 8 deletions.
diff --git a/examples/layers/test_w8a16_gemm.py b/examples/layers/test_w8a16_gemm.py
@@ -1,11 +1,7 @@
-import psutil
-import os
 import time
 import torch
-import torch.nn as nn
 import random
 import numpy as np
-import math
 
 from EETQ import quant_weights, preprocess_weights, w8_a16_gemm
 

diff --git a/examples/models/llama_transformers_example.py b/examples/models/llama_transformers_example.py
@@ -1,6 +1,5 @@
 import os
 import time
-import psutil
 import random
 import torch
 import numpy as np

diff --git a/setup.py b/setup.py
@@ -14,11 +14,11 @@
 from torch.utils import cpp_extension
 from torch.utils.cpp_extension import BuildExtension, CUDAExtension, CppExtension, CUDA_HOME
 
-__version__ = "1.0.0"
+__version__ = "1.0.1"
 
 # ninja build does not work unless include_dirs are abs path
 current_dir = os.path.dirname(os.path.abspath(__file__))
-
+TORCH_VERSION = str(os.getenv("TORCH_VERSION", None) or torch.__version__).split('+', maxsplit=1)[0]
 
 def read(filename):
     with io.open(filename, mode="r", encoding='utf-8') as fd:
@@ -137,7 +137,7 @@ def check_cuda_torch_binary_vs_bare_metal(cuda_dir):
     cmdclass={"build_ext": BuildExtension} if ext_modules else {},
     python_requires=">=3.7",
     install_requires=[
-        "torch>=2.0.1",
+        f"torch>={TORCH_VERSION}",
         "einops",
         "packaging",
         "ninja",