updating cuda and rocm to pytorch 2.4 spec

Docker/Dockerfile

@@ -72,7 +72,7 @@ RUN apt-get update && apt-get install -y --no-install-recommends \
     rm -rf /var/lib/apt/lists/* /tmp/* /var/tmp/*
 
 ARG PYTHON_VERSION=3.10
-ARG FORGE_VERSION=23.11.0-0
+ARG FORGE_VERSION=24.3.0-0
 
 # Install conda
 RUN wget --no-check-certificate -qO ~/miniforge.sh \

Docker/build.py

@@ -30,9 +30,8 @@
 Target = Literal['runtime', 'build_common', 'build_conda', 'build_freesurfer',
                  'build_base', 'runtime_cuda']
 CacheType = Literal["inline", "registry", "local", "gha", "s3", "azblob"]
-AllDeviceType = Literal["cpu", "cuda", "cu116", "cu117", "cu118", "rocm", "rocm5.1.1",
-                        "rocm5.4.2"]
-DeviceType = Literal["cpu", "cu116", "cu117", "cu118", "rocm5.1.1", "rocm5.4.2"]
+AllDeviceType = Literal["cpu", "cuda", "cu118", "cu121", "cu124", "rocm6.1"]
+DeviceType = Literal["cpu", "cu118", "cu121", "cu124", "rocm6.1"]
 
 CREATE_BUILDER = "Create builder with 'docker buildx create --name fastsurfer'."
 CONTAINERD_MESSAGE = (
@@ -61,8 +60,8 @@ class DEFAULTS:
     # torch 2.0.1 comes compiled with cu117, cu118, and rocm5.4.2
     MapDeviceType: Dict[AllDeviceType, DeviceType] = dict(
         ((d, d) for d in get_args(DeviceType)),
-        rocm="rocm5.1.1",
-        cuda="cu117",
+        rocm="rocm6.1",
+        cuda="cu124",
     )
     BUILD_BASE_IMAGE = "ubuntu:22.04"
     RUNTIME_BASE_IMAGE = "ubuntu:22.04"
@@ -185,12 +184,12 @@ def make_parser() -> argparse.ArgumentParser:
 
     parser.add_argument(
         "--device",
-        choices=["cpu", "cuda", "cu117", "cu118", "rocm", "rocm5.4.2"],
+        choices=["cpu", "cuda", "cu118", "cu121", "cu124", "rocm", "rocm6.1"],
         required=True,
         help="""selection of internal build stages to build for a specific platform.<br>
-                - cuda: defaults to cu118, cuda 11.8<br>
+                - cuda: defaults to cu124, cuda 12.4<br>
                 - cpu: only cpu support<br>
-                - rocm: defaults to rocm5.4.2 (experimental)""",
+                - rocm: defaults to rocm6.1 (experimental)""",
     )
     parser.add_argument(
         "--tag",

env/fastsurfer.yml

@@ -5,28 +5,28 @@ channels:
   - defaults
 
 dependencies:
-  - h5py=3.7.0
-  - lapy=1.0.1
-  - matplotlib=3.7.1
-  - nibabel=5.1.0
-  - numpy=1.25.0
-  - pandas=1.5.3
-  - pillow=10.0.1
-  - pip=23.1.2
-  - python=3.10
-  - python-dateutil=2.8.2
-  - pyyaml=6.0
-  - scikit-image=0.19.3
-  - scikit-learn=1.2.2
-  - scipy=1.10.1
-  - setuptools=67.8.0
-  - tensorboard=2.12.1
-  - tqdm=4.66
-  - yacs=0.1.8
-  - pip
-  - pip:
-    - --extra-index-url https://download.pytorch.org/whl/cu117
-    - simpleitk==2.2.1
-    - torch==2.0.1
-    - torchio==0.18.83
-    - torchvision==0.15.2
+- h5py=3.11.0
+- lapy=1.1.0
+- matplotlib=3.9.2
+- nibabel=5.2.1
+- numpy=1.26.4
+- pandas=2.2.2
+- pillow=10.4.0
+- pip=24.2
+- python=3.10
+- python-dateutil=2.9.0
+- pyyaml=6.0.2
+- requests=2.32.3
+- scikit-image=0.24.0
+- scikit-learn=1.5.1
+- scipy=1.14.1
+- setuptools=72.2.0
+- tensorboard=2.17.1
+- tqdm=4.66.5
+- yacs=0.1.8
+- pip:
+  - --extra-index-url https://download.pytorch.org/whl/cu124
+  - simpleitk==2.4.0
+  - torch==2.4.0+cu124
+  - torchio==0.19.9
+  - torchvision==0.19.0+cu124