DiNTS bundles: make search code loadable with torch.load(weights_only=True)

models/multi_organ_segmentation/configs/inference.yaml

@@ -5,7 +5,7 @@ imports:
 input_channels: 1
 output_classes: 8
 arch_ckpt_path: "$@bundle_root + '/models/search_code_18590.pt'"
-arch_ckpt: "$torch.load(@arch_ckpt_path, map_location=torch.device('cuda'))"
+arch_ckpt: "$torch.load(@arch_ckpt_path, map_location=@device)"
 bundle_root: "."
 output_dir: "$@bundle_root + '/eval'"
 dataset_dir: "/data/"
@@ -20,16 +20,16 @@ dints_space:
   num_depths: 4
   use_downsample: true
   arch_code:
-  - "$@arch_ckpt['arch_code_a']"
-  - "$@arch_ckpt['arch_code_c']"
-  device: "$torch.device('cuda')"
+  - "$@arch_ckpt['arch_code_a'].cpu().numpy()"
+  - "$@arch_ckpt['arch_code_c'].cpu().numpy()"
+  device: "@device"
 network_def:
   _target_: monai.networks.nets.DiNTS
   dints_space: "@dints_space"
   in_channels: "@input_channels"
   num_classes: "@output_classes"
   use_downsample: true
-  node_a: "$torch.from_numpy(@arch_ckpt['node_a'])"
+  node_a: "$torch.as_tensor(@arch_ckpt['node_a'])"
 network: "$@network_def.to(@device)"
 preprocessing:
   _target_: Compose

models/multi_organ_segmentation/configs/metadata.json

@@ -1,7 +1,8 @@
 {
     "schema": "https://github.com/Project-MONAI/MONAI-extra-test-data/releases/download/0.8.1/meta_schema_20240725.json",
-    "version": "0.0.6",
+    "version": "0.0.7",
     "changelog": {
+        "0.0.7": "ship search code in the bundle as a tensor checkpoint loadable with torch.load(weights_only=True); search.py saves tensors; use the configured device instead of hard-coded cuda",
         "0.0.6": "enhance metadata with improved descriptions",
         "0.0.5": "update to huggingface hosting",
         "0.0.4": "Set image_only to False",

models/multi_organ_segmentation/configs/train.yaml

@@ -8,7 +8,7 @@ imports:
 input_channels: 1
 output_classes: 8
 arch_ckpt_path: "$@bundle_root + '/models/search_code_18590.pt'"
-arch_ckpt: "$torch.load(@arch_ckpt_path, map_location=torch.device('cuda'))"
+arch_ckpt: "$torch.load(@arch_ckpt_path, map_location=@device)"
 bundle_root: "."
 ckpt_dir: "$@bundle_root + '/models'"
 output_dir: "$@bundle_root + '/eval'"
@@ -26,9 +26,9 @@ dints_space:
   num_depths: 4
   use_downsample: true
   arch_code:
-  - "$@arch_ckpt['arch_code_a']"
-  - "$@arch_ckpt['arch_code_c']"
-  device: "$torch.device('cuda')"
+  - "$@arch_ckpt['arch_code_a'].cpu().numpy()"
+  - "$@arch_ckpt['arch_code_c'].cpu().numpy()"
+  device: "@device"
 network_def:
   _target_: monai.networks.nets.DiNTS
   dints_space: "@dints_space"

models/multi_organ_segmentation/large_files.yml

@@ -1,8 +1,4 @@
 large_files:
-  - path: "models/search_code_18590.pt"
-    url: "https://drive.google.com/uc?export=download&id=1G2YZo1HKkvf5sMvaaRQeZp_YA9GDlvL-"
-    hash_val: "55dc7fe4bee4a93d25a1dd329dfdd159"
-    hash_type: "md5"
   - path: "models/model.pt"
     url: "https://drive.google.com/uc?export=download&id=1kH0yTyiXUNqdYXpnSXI2p5-vFwDYoCzl"
     hash_val: "e21d663b33d29a8ca4e28d60a14a3d66"

models/pancreas_ct_dints_segmentation/configs/inference.yaml

@@ -6,7 +6,7 @@ imports:
 input_channels: 1
 output_classes: 3
 arch_ckpt_path: "$@bundle_root + '/models/search_code_18590.pt'"
-arch_ckpt: "$torch.load(@arch_ckpt_path, map_location=torch.device('cuda'))"
+arch_ckpt: "$torch.load(@arch_ckpt_path, map_location=@device)"
 bundle_root: "."
 image_key: "image"
 output_dir: "$@bundle_root + '/eval'"
@@ -27,16 +27,16 @@ dints_space:
   num_depths: 4
   use_downsample: true
   arch_code:
-  - "$@arch_ckpt['arch_code_a']"
-  - "$@arch_ckpt['arch_code_c']"
-  device: "$torch.device('cuda')"
+  - "$@arch_ckpt['arch_code_a'].cpu().numpy()"
+  - "$@arch_ckpt['arch_code_c'].cpu().numpy()"
+  device: "@device"
 network_def:
   _target_: monai.networks.nets.DiNTS
   dints_space: "@dints_space"
   in_channels: "@input_channels"
   num_classes: "@output_classes"
   use_downsample: true
-  node_a: "$torch.from_numpy(@arch_ckpt['node_a'])"
+  node_a: "$torch.as_tensor(@arch_ckpt['node_a'])"
 network: "$@network_def.to(@device)"
 preprocessing:
   _target_: Compose

models/pancreas_ct_dints_segmentation/configs/metadata.json

@@ -1,7 +1,8 @@
 {
     "schema": "https://github.com/Project-MONAI/MONAI-extra-test-data/releases/download/0.8.1/meta_schema_20240725.json",
-    "version": "0.5.3",
+    "version": "0.5.4",
     "changelog": {
+        "0.5.4": "ship search code in the bundle as a tensor checkpoint loadable with torch.load(weights_only=True); search.py saves tensors; use the configured device instead of hard-coded cuda",
         "0.5.3": "modernize type annotations and string formatting without changing behavior",
         "0.5.2": "enhance metadata with improved descriptions",
         "0.5.1": "update to huggingface hosting",

models/pancreas_ct_dints_segmentation/configs/train.yaml

@@ -8,7 +8,7 @@ imports:
 input_channels: 1
 output_classes: 3
 arch_ckpt_path: "$@bundle_root + '/models/search_code_18590.pt'"
-arch_ckpt: "$torch.load(@arch_ckpt_path, map_location=torch.device('cuda'))"
+arch_ckpt: "$torch.load(@arch_ckpt_path, map_location=@device)"
 bundle_root: "."
 ckpt_dir: "$@bundle_root + '/models'"
 output_dir: "$@bundle_root + '/eval'"
@@ -26,9 +26,9 @@ dints_space:
   num_depths: 4
   use_downsample: true
   arch_code:
-  - "$@arch_ckpt['arch_code_a']"
-  - "$@arch_ckpt['arch_code_c']"
-  device: "$torch.device('cuda')"
+  - "$@arch_ckpt['arch_code_a'].cpu().numpy()"
+  - "$@arch_ckpt['arch_code_c'].cpu().numpy()"
+  device: "@device"
 network_def:
   _target_: monai.networks.nets.DiNTS
   dints_space: "@dints_space"

models/pancreas_ct_dints_segmentation/large_files.yml

@@ -1,8 +1,4 @@
 large_files:
-- path: "models/search_code_18590.pt"
-  url: "https://developer.download.nvidia.com/assets/Clara/monai/tutorials/model_zoo/model_pancreas_ct_dints_segmentation_search_code.pt"
-  hash_val: "2ff53724aa9d80d9025034d36785cdb8"
-  hash_type: "md5"
 - path: "models/model.pt"
   url: "https://developer.download.nvidia.com/assets/Clara/monai/tutorials/model_zoo/model_pancreas_ct_dints_segmentation.pt"
   hash_val: "ce8531b5143b417f65d5b481f72bc911"

models/pancreas_ct_dints_segmentation/scripts/search.py

@@ -458,10 +458,12 @@ def run(config_file: str | Sequence[str]):  # ruff: ignore[C901]  FIXME: too com
                     node_a_d, arch_code_a_d, arch_code_c_d, arch_code_a_max_d = dints_space.decode()
                     torch.save(
                         {
-                            "node_a": node_a_d,
-                            "arch_code_a": arch_code_a_d,
-                            "arch_code_a_max": arch_code_a_max_d,
-                            "arch_code_c": arch_code_c_d,
+                            # save as tensors (not numpy arrays) so the file can be read with
+                            # torch.load(weights_only=True), the default since PyTorch 2.6
+                            "node_a": torch.as_tensor(node_a_d),
+                            "arch_code_a": torch.as_tensor(arch_code_a_d),
+                            "arch_code_a_max": torch.as_tensor(arch_code_a_max_d),
+                            "arch_code_c": torch.as_tensor(arch_code_c_d),
                             "iter_num": idx_iter,
                             "epochs": epoch + 1,
                             "best_dsc": best_metric,