Accept torch tensors as DiNTS arch_code for weights_only=True checkpoints

TopologyConstruction.__init__ converted arch_code entries with
torch.from_numpy, which raises TypeError when the architecture codes are
already tensors. Checkpoints such as search_code_18590.pt from the
pancreas_ct_dints_segmentation and multi_organ_segmentation bundles must
be re-saved with tensors to load under torch.load(weights_only=True),
so the network needs to accept both. torch.as_tensor handles numpy
arrays and tensors identically, keeping the existing numpy path intact.

Fixes #9025 (MONAI code side; bundle artifact re-upload tracked separately)

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Signed-off-by: Minsu Kim <minsu.kim08@gmail.com>

Author: minsuking

monai/networks/nets/dints.py

@@ -519,7 +519,8 @@ class TopologyConstruction(nn.Module):
     The base class for `TopologyInstance` and `TopologySearch`.
 
     Args:
-        arch_code: `[arch_code_a, arch_code_c]`, numpy arrays. The architecture codes defining the model.
+        arch_code: `[arch_code_a, arch_code_c]`, numpy arrays or torch tensors. The architecture codes defining
+            the model.
             For example, for a ``num_depths=4, num_blocks=12`` search space:
 
             - `arch_code_a` is a 12x10 (10 paths) binary matrix representing if a path is activated.
@@ -605,8 +606,8 @@ def __init__(
             arch_code_a = torch.ones((self.num_blocks, len(self.arch_code2out))).to(self.device)
             arch_code_c = torch.ones((self.num_blocks, len(self.arch_code2out), self.num_cell_ops)).to(self.device)
         else:
-            arch_code_a = torch.from_numpy(arch_code[0]).to(self.device)
-            arch_code_c = F.one_hot(torch.from_numpy(arch_code[1]).to(torch.int64), self.num_cell_ops).to(self.device)
+            arch_code_a = torch.as_tensor(arch_code[0]).to(self.device)
+            arch_code_c = F.one_hot(torch.as_tensor(arch_code[1]).to(torch.int64), self.num_cell_ops).to(self.device)
 
         self.arch_code_a = arch_code_a
         self.arch_code_c = arch_code_c

tests/networks/nets/test_dints_network.py

@@ -153,6 +153,48 @@ def test_dints_search(self, dints_grid_params, dints_params, input_shape, expect
         self.assertTrue(isinstance(net.weight_parameters(), list))
 
 
+class TestDintsArchCode(unittest.TestCase):
+    """arch_code entries loaded with `torch.load(..., weights_only=True)` are tensors, not numpy arrays."""
+
+    def setUp(self):
+        self.grid_params = {
+            "channel_mul": 0.2,
+            "num_blocks": 6,
+            "num_depths": 3,
+            "device": "cpu",
+            "use_downsample": False,
+            "spatial_dims": 3,
+        }
+        _cell = Cell(1, 1, 0, spatial_dims=self.grid_params["spatial_dims"])
+        num_paths = 3 * self.grid_params["num_depths"] - 2
+        self.arch_code_a = np.ones((self.grid_params["num_blocks"], num_paths))
+        self.arch_code_c = np.random.randint(len(_cell.OPS), size=(self.grid_params["num_blocks"], num_paths))
+
+    def test_tensor_arch_code_matches_numpy(self):
+        grid_np = TopologyInstance(arch_code=[self.arch_code_a, self.arch_code_c], **self.grid_params)
+        grid_pt = TopologyInstance(
+            arch_code=[torch.as_tensor(self.arch_code_a), torch.as_tensor(self.arch_code_c)], **self.grid_params
+        )
+        torch.testing.assert_close(grid_pt.arch_code_a, grid_np.arch_code_a)
+        torch.testing.assert_close(grid_pt.arch_code_c, grid_np.arch_code_c)
+        self.assertEqual(set(grid_pt.cell_tree.keys()), set(grid_np.cell_tree.keys()))
+
+    def test_dints_forward_tensor_arch_code(self):
+        grid = TopologyInstance(
+            arch_code=[torch.as_tensor(self.arch_code_a), torch.as_tensor(self.arch_code_c)], **self.grid_params
+        )
+        net = DiNTS(
+            dints_space=grid,
+            in_channels=1,
+            num_classes=2,
+            spatial_dims=3,
+            use_downsample=False,
+            node_a=torch.ones((self.grid_params["num_blocks"] + 1, self.grid_params["num_depths"])),
+        )
+        result = net(torch.randn(1, 1, 16, 16, 16))
+        self.assertEqual(result.shape, (1, 2, 16, 16, 16))
+
+
 class TestDintsTS(unittest.TestCase):
     @parameterized.expand(TEST_CASES_3D + TEST_CASES_2D)
     def test_script(self, dints_grid_params, dints_params, input_shape, _):