Skip to content

Commit b01550f

Browse files
committed
Merge branch 'bitpacking' of https://github.com/vayuda/ao into bitpacking
2 parents 4ce1295 + f6758e9 commit b01550f

File tree

1 file changed

+3
-2
lines changed

1 file changed

+3
-2
lines changed

setup.py

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@ def read_requirements(file_path):
1818

1919
# Determine the package name based on the presence of an environment variable
2020
package_name = "torchao-nightly" if os.environ.get("TORCHAO_NIGHTLY") else "torchao"
21+
version_suffix = os.getenv("VERSION_SUFFIX", "")
2122

2223
# Version is year.month.date if using nightlies
2324
version = current_date if package_name == "torchao-nightly" else "0.2.0"
@@ -42,7 +43,7 @@ def get_extensions():
4243
if CUDA_HOME is None and torch.cuda.is_available():
4344
print("CUDA toolkit is not available. Skipping compilation of CUDA extensions")
4445
print("If you'd like to compile CUDA extensions locally please install the cudatoolkit from https://anaconda.org/nvidia/cuda-toolkit")
45-
46+
4647
use_cuda = torch.cuda.is_available() and CUDA_HOME is not None
4748
extension = CUDAExtension if use_cuda else CppExtension
4849

@@ -85,7 +86,7 @@ def get_extensions():
8586

8687
setup(
8788
name=package_name,
88-
version=version,
89+
version=version+version_suffix,
8990
packages=find_packages(),
9091
include_package_data=True,
9192
package_data={

0 commit comments

Comments
 (0)