Support specifying a torch range

pytorch · Feb 25, 2025 · 4272b38 · 4272b38
1 parent c670ad8
commit 4272b38
Showing 1 changed file with 6 additions and 3 deletions.
diff --git a/setup.py b/setup.py
@@ -39,9 +39,12 @@ def _make_version_file(version, sha):
 
 
 def _get_pytorch_version():
-    if "PYTORCH_VERSION" in os.environ:
-        return f"torch=={os.environ['PYTORCH_VERSION']}"
-    return "torch"
+    pytorch_dep = os.getenv("TORCH_PACKAGE_NAME", "torch")
+    if version_pin := os.getenv("PYTORCH_VERSION"):
+        pytorch_dep += "==" + version_pin
+    elif (version_pin_ge := os.getenv("PYTORCH_VERSION_GE")) and (version_pin_lt := os.getenv("PYTORCH_VERSION_LT")):
+        pytorch_dep += f">={version_pin_ge},<{version_pin_lt}"
+    return pytorch_dep
 
 
 class clean(distutils.command.clean.clean):