fix: torchvision version in environment.yaml file

src/careamics/model_io/bioimage/bioimage_utils.py

@@ -22,7 +22,7 @@ def get_unzip_path(zip_path: Union[Path, str]) -> Path:
     return zip_path.parent / (str(zip_path.name) + ".unzip")
 
 
-def create_env_text(pytorch_version: str) -> str:
+def create_env_text(pytorch_version: str, torchvision_version: str) -> str:
     """Create environment yaml content for the bioimage model.
 
     This installs an environment with the specified pytorch version and the latest
@@ -32,6 +32,8 @@ def create_env_text(pytorch_version: str) -> str:
     ----------
     pytorch_version : str
         Pytorch version.
+    torchvision_version : str
+        Torchvision version.
 
     Returns
     -------
@@ -43,7 +45,7 @@ def create_env_text(pytorch_version: str) -> str:
         f"dependencies:\n"
         f"  - python=3.10\n"
         f"  - pytorch={pytorch_version}\n"
-        f"  - torchvision={pytorch_version}\n"
+        f"  - torchvision={torchvision_version}\n"
         f"  - pip\n"
         f"  - pip:\n"
         f"    - git+https://github.com/CAREamics/careamics.git\n"

src/careamics/model_io/bmz_io.py

@@ -8,7 +8,9 @@
 import pkg_resources
 from bioimageio.core import load_description, test_model
 from bioimageio.spec import ValidationSummary, save_bioimageio_package
-from torch import __version__, load, save
+from torch import __version__ as PYTORCH_VERSION
+from torch import load, save
+from torchvision import __version__ as TORCHVISION_VERSION
 
 from careamics.config import Configuration, load_configuration, save_configuration
 from careamics.config.support import SupportedArchitecture
@@ -141,7 +143,6 @@ def export_to_bmz(
         path_to_archive.parent.mkdir(parents=True, exist_ok=True)
 
     # versions
-    pytorch_version = __version__
     careamics_version = pkg_resources.get_distribution("careamics").version
 
     # save files in temporary folder
@@ -151,7 +152,7 @@ def export_to_bmz(
         # create environment file
         # TODO move in bioimage module
         env_path = temp_path / "environment.yml"
-        env_path.write_text(create_env_text(pytorch_version))
+        env_path.write_text(create_env_text(PYTORCH_VERSION, TORCHVISION_VERSION))
 
         # export input and ouputs
         inputs = temp_path / "inputs.npy"
@@ -174,7 +175,7 @@ def export_to_bmz(
             inputs=inputs,
             outputs=outputs,
             weights_path=weight_path,
-            torch_version=pytorch_version,
+            torch_version=PYTORCH_VERSION,
             careamics_version=careamics_version,
             config_path=config_path,
             env_path=env_path,