pin numpy for pytorch env

bioimage-io · Jul 9, 2024 · 10d246c · 10d246c
1 parent 67c17ca
commit 10d246c
Showing 1 changed file with 10 additions and 0 deletions.
diff --git a/bioimageio_collection_backoffice/conda_env.py b/bioimageio_collection_backoffice/conda_env.py
@@ -139,6 +139,7 @@ def _get_default_pytorch_env(
 
     # avoid `undefined symbol: iJIT_NotifyEvent` from `torch/lib/libtorch_cpu.so`
     # see https://github.com/pytorch/pytorch/issues/123097
+    # TODO: check if this is fixed for future pytorch versions
     deps.append("mkl ==2024.0.0")
 
     if pytorch_version.major == 1 or (
@@ -148,6 +149,15 @@ def _get_default_pytorch_env(
         # see https://github.com/pypa/setuptools/issues/4376#issuecomment-2126162839
         deps.append("setuptools <70.0.0")
 
+    if pytorch_version < Version(
+        "2.4"
+    ):  # TODO: verify that future pytorch 2.4 is numpy 2.0 compatible
+        # make sure we have a compatible numpy version
+        # see https://github.com/pytorch/vision/issues/8460
+        deps.append("numpy <2")
+    else:
+        deps.append("numpy >=2,<3")
+
     return CondaEnv(
         name="env",
         channels=["pytorch", "conda-forge"],