require torch <2.6.0

wasserth · Jan 30, 2025 · 611b54a · 611b54a
1 parent 6b226fc
commit 611b54a
Show file tree

Hide file tree

Showing 3 changed files with 4 additions and 4 deletions.
diff --git a/README.md b/README.md
@@ -19,7 +19,7 @@ TotalSegmentator works on Ubuntu, Mac, and Windows and on CPU and GPU.
 
 Install dependencies:
 * Python >= 3.9
-* [Pytorch](http://pytorch.org/) >= 2.0.0 (and <2.4 for windows)
+* [Pytorch](http://pytorch.org/) >= 2.0.0 and <2.6.0 (and <2.4 for windows)
 
 Optionally:
 * if you use the option `--preview` you have to install xvfb (`apt-get install xvfb`) and fury (`pip install fury`)

diff --git a/setup.py b/setup.py
@@ -18,7 +18,7 @@
              "resources/ct_brain_atlas_1mm.nii.gz"]
             },
         install_requires=[
-            'torch>=2.1.2',
+            'torch>=2.1.2,<2.6.0',
             'numpy<2',
             'SimpleITK',
             'nibabel>=2.3.0',

diff --git a/totalsegmentator/bin/totalseg_evans_index.py b/totalsegmentator/bin/totalseg_evans_index.py
@@ -28,12 +28,12 @@ def run_models(ct_img, verbose=False):
     ventricle_parts = totalsegmentator(ct_img, None, task="ventricle_parts", ml=True, nr_thr_saving=1, quiet=not verbose)
     return brain_skull, ventricle_parts
 
-
+# Required if calling from e.g. streamlit; calling python_api not working there
 def run_models_shell(ct_img, verbose=False):
     quiet = "--quiet" if not verbose else ""
     with tempfile.TemporaryDirectory(prefix="totalseg_tmp_2_") as tmp_folder:
         tmp_dir = Path(tmp_folder)
-        ct_img_path = tmp_dir / "ct.nii.gz"  # save to tmp folder
+        ct_img_path = tmp_dir / "ct.nii.gz"
         nib.save(ct_img, ct_img_path)
         subprocess.call(f"TotalSegmentator -i {ct_img_path} -o {tmp_dir / 'brain_skull.nii.gz'} --roi_subset brain skull --ml --nr_thr_saving 1 {quiet}", shell=True)
         subprocess.call(f"TotalSegmentator -i {ct_img_path} -o {tmp_dir / 'ventricle_parts.nii.gz'} --task ventricle_parts --ml --nr_thr_saving 1 {quiet}", shell=True)