Accept torch tensors as DiNTS arch_code for weights_only=True checkpoints

[minsuking]

Addresses the MONAI-side portion of #9025. This PR does not fully resolve that issue
on its own — see "Follow-up required" below for the bundle-artifact half.

Description

TopologyConstruction.__init__ converted the arch_code entries with
torch.from_numpy(...), which raises TypeError: expected np.ndarray (got Tensor)
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 currently store numpy arrays, which
torch.load(..., weights_only=True) rejects; once those files are re-saved with
tensors, the DiNTS network must accept tensor arch codes to load them.

This PR replaces the two torch.from_numpy(...) calls with torch.as_tensor(...)
(monai/networks/nets/dints.py), which handles numpy arrays and tensors
identically:

• numpy path unchanged: torch.as_tensor(ndarray) is equivalent to
  torch.from_numpy — zero-copy (same data_ptr()), dtype preserved, followed by
  the same .to(self.device).
• tensor path: torch.as_tensor(tensor) returns the input as-is (no copy,
  no warning), unlike torch.tensor(...) which always copies and warns.
• arch_code_c still goes through .to(torch.int64) before F.one_hot, so the
  one-hot output dtype (int64) is identical for both input types; all downstream
  consumers (== 1 gates, arch_c > 0 in MixedOp, forward-path activation
  checks) already operated on tensors after the conversion site, so their behavior
  is unchanged.

Credit where due: another contributor (@Theerth1) independently traced the same
root cause to TopologyConstruction.__init__ in the issue discussion before this
PR was prepared.

Follow-up required (not resolved by this PR)

This change alone does not make the affected bundles loadable with
weights_only=True: the shipped search_code_18590.pt files contain pickled
numpy arrays, which the weights_only unpickler rejects at torch.load time
regardless of this fix. The checkpoints for pancreas_ct_dints_segmentation and
multi_organ_segmentation (including their identical HuggingFace mirrors) need to
be re-saved with tensors and re-uploaded in Project-MONAI/model-zoo. This PR is
the prerequisite that lets those re-saved tensor checkpoints construct the
network. Per the investigation in #9025, no other model-zoo bundles are affected.

Notes / disclosed risks

• torch.as_tensor also accepts Python lists/scalars, so the accepted input is
  slightly broader than the documented numpy arrays or torch tensors; this is
  benign.
• arch_code_a's stored dtype now follows a tensor input's dtype (previously
  always the numpy dtype); all consumers are dtype-agnostic comparisons, and no
  functional difference was found.
• Verified locally (Windows, Python 3.11): full
  tests/networks/nets/test_dints_network.py module passes via unittest
  (including the TorchScript tests); the new TestDintsArchCode cases pass;
  ruff check clean on both changed files; mypy monai/networks/nets/dints.py
  clean. black was not available locally, so formatting is deferred to CI.

Types of changes

• Non-breaking change (fix or new feature that would not break existing functionality).
• New tests added to cover the changes.
• In-line docstrings updated.

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

• Configuration used: Repository: Project-MONAI/MONAI/.coderabbit.yaml
• Review profile: CHILL
• Plan: Advanced
• Run ID: 7dbcdc3b-092f-4760-9fd7-0ea57afec56d

📥 Commits

Reviewing files that changed from the base of the PR and between 30f8875 and 5a826e4.

📒 Files selected for processing (2)

• monai/networks/nets/dints.py
• tests/networks/nets/test_dints_network.py

Included review availability: This review used your included allowance. Your plan provides up to 8 included reviews per hour; 7 remain after this review.

---

📝 Walkthrough

Walkthrough

TopologyConstruction accepts NumPy arrays and PyTorch tensors for arch_code. It converts the inputs with torch.as_tensor. New CPU tests compare the resulting topologies and check DiNTS output shape with tensor codes.

Priority: ➖ Normal

Estimated code review effort: 2 (Simple) | ~10 minutes

Merge Risk: ⚪ Minimal · up to 5a826

The DiNTS tensor-code change is ready for normal merge checks. Re-saving and publishing the affected checkpoints remains separate work.

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 50.00% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 6 functions across 2 files.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly identifies the DiNTS change to accept tensor architecture codes. Its reference to weights_only checkpoints provides relevant context, although the PR is only a prerequisite for loadi…
Description check	✅ Passed	The description explains the change, its rationale, tests, and the required bundle follow-up. It omits the template's “Fixes #” line and leaves the integration and quick-test checkboxes unmarked, but …
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

• Fix all pre-merge checks with AI

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create a new PR

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}

[coderabbitai]

🧹 Nitpick comments (1)

tests/networks/nets/test_dints_network.py (1)

159-171: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Add docstrings to the new test routines.

setUp, test_tensor_arch_code_matches_numpy, and test_dints_forward_tensor_arch_code are new definitions without docstrings. Add concise Google-style docstrings that describe the fixture data, input representations, and expected output. Use Args, Returns, and Raises sections when applicable.

As per path instructions: “Docstrings should be present for all definition which describe each variable, return value, and raised exception in the appropriate section of the Google-style of docstrings.”

Also applies to: 173-195

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@tests/networks/nets/test_dints_network.py` around lines 159 - 171, Add
concise Google-style docstrings to setUp, test_tensor_arch_code_matches_numpy,
and test_dints_forward_tensor_arch_code, documenting the fixture data,
tensor/NumPy architecture-code inputs, expected outputs, and any applicable
arguments or exceptions using Args, Returns, and Raises sections. Keep the test
behavior unchanged.

Source: Path instructions

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Nitpick comments:
In `@tests/networks/nets/test_dints_network.py`:
- Around line 159-171: Add concise Google-style docstrings to setUp,
test_tensor_arch_code_matches_numpy, and test_dints_forward_tensor_arch_code,
documenting the fixture data, tensor/NumPy architecture-code inputs, expected
outputs, and any applicable arguments or exceptions using Args, Returns, and
Raises sections. Keep the test behavior unchanged.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro Plus

Run ID: 6bcce4cc-865d-4087-b5c0-784b6a861cd1

📥 Commits

Reviewing files that changed from the base of the PR and between 87060c4 and 54e7bee.

📒 Files selected for processing (2)

• monai/networks/nets/dints.py
• tests/networks/nets/test_dints_network.py

[ericspod]

Hi @minsuking thanks for this fix, we do need to revisit saved bundle files to resave them as tensors.