Fix AutoencoderKlMaisi forcing CUDA transfer on CPU inputs

[ytl0623]

Fixes #8735

Description

This PR fixes a bug in AutoencoderKlMaisi where the model would force a transfer to cuda even if the input tensors and the model were placed on the CPU.

Types of changes

• Non-breaking change (fix or new feature that would not break existing functionality).
• Breaking change (fix or new feature that would cause existing functionality to change).
• New tests added to cover the changes.
• Integration tests passed locally by running ./runtests.sh -f -u --net --coverage.
• Quick tests passed locally by running ./runtests.sh --quick --unittests --disttests.
• In-line docstrings updated.
• Documentation updated, tested make html command in the docs/ folder.

[coderabbitai]

📝 Walkthrough

Walkthrough

The _concatenate_tensors method in MaisiConvolution now captures the original device of the first tensor before concatenation and conditionally returns the result to that device. If the original device is CPU, the tensor remains on CPU; if CUDA, it moves back to CUDA. This replaces the previous unconditional transfer to CUDA, preventing unnecessary device transfers for CPU-based inputs.

Estimated code review effort

🎯 2 (Simple) | ⏱️ ~8 minutes

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 0.00% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Title check	✅ Passed	Title clearly summarizes the main fix: preventing forced CUDA transfer on CPU inputs in AutoencoderKlMaisi.
Description check	✅ Passed	Description covers the bug fix with issue reference, clear explanation of the problem, and appropriate checkboxes marked.
Linked Issues check	✅ Passed	Code changes directly address issue #8735 by making CUDA transfer conditional on the original device instead of unconditional.
Out of Scope Changes check	✅ Passed	All changes are focused solely on fixing the device transfer bug in _concatenate_tensors, with no unrelated modifications.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing touches

• 📝 Generate docstrings

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Post copyable unit tests in a comment

---

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 and usage tips.}

[coderabbitai]

Caution

Some comments are outside the diff and can’t be posted inline due to platform limitations.

⚠️ Outside diff range comments (1)

monai/apps/generation/maisi/networks/autoencoderkl_maisi.py (1)

217-232: ⚠️ Potential issue | 🟡 Minor

Fix is correct and improves on the device-handling pattern.

Unconditionally moves tensors to CPU for concatenation (line 219), then conditionally returns to original device only if non-CPU (line 230-231). Unlike _cat_inputs which hardcodes "cuda", this approach works for any device type.

However, existing tests use input shape (1, 1, 32, 32, 32) with max size 32, which falls below the 500-element threshold (line 214) and skips this large-tensor concatenation branch entirely. A test with max(size) ≥ 500 is needed to exercise and prevent regression in lines 217-231.

[ericspod]

Thanks for the fix @ytl0623, looks good to me.