Skip to content

Commit 267f0e1

Browse files
author
Pablo Gonzalez
committed
fix problem with Pytorch versions
Signed-off-by: Pablo Gonzalez <[email protected]>
1 parent eee6ea3 commit 267f0e1

File tree

1 file changed

+4
-5
lines changed

1 file changed

+4
-5
lines changed

setup.py

Lines changed: 4 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -12,21 +12,20 @@
1212
from skbuild import setup
1313

1414
INSTALL_REQUIRES = [
15-
"torch{}".format(os.getenv("TORCH_VERSION_SPECIFIER", ">=1.9")),
15+
"torch{}".format(os.getenv("TORCH_VERSION_SPECIFIER", ">=1.9,<2.7")),
1616
"torchvision",
1717
"scipy",
1818
"requests>=2.25,<3",
1919
"numpy>=1.22",
2020
"protobuf>=4.21.6",
2121
]
2222

23-
2423
def get_version() -> str:
25-
"""Get the package version."""
24+
"""Get the package version for CUDA enabled wheels."""
2625
version_path = os.path.join(os.path.dirname(__file__), "src", "aihwkit", "VERSION.txt")
2726
with open(version_path, encoding="utf-8") as version_file:
28-
return version_file.read().strip()
29-
27+
v = version_file.read().strip()
28+
return f"{v}{os.getenv('AIHWKIT_VERSION_SUFFIX', '')}"
3029

3130
def get_long_description() -> str:
3231
"""Get the package long description."""

0 commit comments

Comments
 (0)