fix SpecifierSet behavior with prereleases (#2959)

By default, SpecifierSet doesn't allow prereleases, which is not our expected behavior. Signed-off-by: Jinzhe Zeng <[email protected]>
deepmodeling · Oct 30, 2023 · 1429e46 · 1429e46
1 parent 9a4a4dd
commit 1429e46
Showing 1 changed file with 5 additions and 5 deletions.
diff --git a/backend/find_tensorflow.py b/backend/find_tensorflow.py
@@ -114,9 +114,9 @@ def get_tf_requirement(tf_version: str = "") -> dict:
 
     extra_requires = []
     extra_select = {}
-    if not (tf_version == "" or tf_version in SpecifierSet(">=2.12")):
+    if not (tf_version == "" or tf_version in SpecifierSet(">=2.12", prereleases=True)):
         extra_requires.append("protobuf<3.20")
-    if tf_version == "" or tf_version in SpecifierSet(">=1.15"):
+    if tf_version == "" or tf_version in SpecifierSet(">=1.15", prereleases=True):
         extra_select["mpi"] = [
             "horovod",
             "mpi4py",
@@ -138,9 +138,9 @@ def get_tf_requirement(tf_version: str = "") -> dict:
             ],
             **extra_select,
         }
-    elif tf_version in SpecifierSet("<1.15") or tf_version in SpecifierSet(
-        ">=2.0,<2.1"
-    ):
+    elif tf_version in SpecifierSet(
+        "<1.15", prereleases=True
+    ) or tf_version in SpecifierSet(">=2.0,<2.1", prereleases=True):
         return {
             "cpu": [
                 f"tensorflow=={tf_version}",