version prefix matching

Signed-off-by: Jinzhe Zeng <[email protected]>
deepmodeling · Jun 22, 2024 · 4c38fad · 4c38fad
1 parent 6a9c294
commit 4c38fad
Showing 1 changed file with 5 additions and 1 deletion.
diff --git a/backend/find_pytorch.py b/backend/find_pytorch.py
@@ -111,7 +111,11 @@ def get_pt_requirement(pt_version: str = "") -> dict:
 
     return {
         "torch": [
-            f"torch=={Version(pt_version).base_version}"
+            # uv has different local version behaviors, i.e. `==2.3.1` cannot match `==2.3.1+cpu`
+            # https://github.com/astral-sh/uv/blob/main/PIP_COMPATIBILITY.md#local-version-identifiers
+            # luckily, .* (prefix matching) defined in PEP 440 can match any local version
+            # https://peps.python.org/pep-0440/#version-matching
+            f"torch=={Version(pt_version).base_version}.*"
             if pt_version != ""
             else "torch>=2a",
         ],