Fix #421: Update UNETR to support proj_type in newer MONAI versions

[aman0311x]

Description
This PR fixes issue #421 where the UNETR model crashes on MONAI >= 1.3.0 due to the deprecation of the pos_embed parameter in the ViT class, which was renamed to proj_type.

Changes Made
Updated the parameter name from pos_embed to proj_type in UNETR.init arguments, docstrings, and examples.

Added a backward compatibility check using Python's inspect module to dynamically pass either proj_type or pos_embed based on the installed MONAI version.

Fixes #421

Summary by CodeRabbit

• Bug Fixes

  • Improved compatibility with supported vision transformer configurations when initializing UNETR.
  • Corrected patch embedding bias loading from pretrained model state data.

• Documentation

  • Updated UNETR configuration guidance and examples to use the revised position embedding setting.

[coderabbitai]

Walkthrough

UNETR now exposes proj_type, validates it, and passes either proj_type or legacy pos_embed depending on the installed ViT signature. The patch embedding bias checkpoint key is also corrected.

Changes

UNETR compatibility updates

Layer / File(s)	Summary
Projection compatibility and checkpoint loading UNETR/BTCV/networks/unetr.py	The constructor documentation, example, validation, and public argument use proj_type; ViT initialization falls back to pos_embed when required, and patch embedding bias loading uses the updated state-dict path.

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

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	Title clearly states the UNETR MONAI compatibility fix and the proj_type rename.
Linked Issues check	✅ Passed	The changes address the linked bug by renaming the ViT arg to proj_type and keeping compatibility with older MONAI versions.
Out of Scope Changes check	✅ Passed	The extra state-dict key adjustment appears tied to the UNETR refactor and does not indicate unrelated scope creep.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

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]

Actionable comments posted: 1

Caution

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

⚠️ Outside diff range comments (1)

UNETR/BTCV/networks/unetr.py (1)

211-215: 🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick win

Checkpoint bias key is missing patch_embedding. prefix — inconsistent with weight key.

Line 212 reads the weight from module.transformer.patch_embedding.patch_embeddings.1.weight, but line 215 reads the bias from module.transformer.patch_embeddings.1.bias. The bias key is missing the patch_embedding. segment, making it inconsistent with the weight key. Since both tensors belong to the same patch_embeddings[1] module in the checkpoint, the bias key should follow the same path prefix.

🐛 Proposed fix: restore the `patch_embedding.` prefix

             self.vit.patch_embedding.patch_embeddings[1].bias.copy_(
-                weights["state_dict"]["module.transformer.patch_embeddings.1.bias"]
+                weights["state_dict"]["module.transformer.patch_embedding.patch_embeddings.1.bias"]
             )

🤖 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 `@UNETR/BTCV/networks/unetr.py` around lines 211 - 215, Update the bias
checkpoint lookup in the weight-loading block to use the same module path as the
corresponding weight lookup, including the patch_embedding segment before
patch_embeddings.1.bias. Keep the target module
self.vit.patch_embedding.patch_embeddings[1] unchanged.

🧹 Nitpick comments (1)

UNETR/BTCV/networks/unetr.py (1)

102-102: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value

Move import inspect to module top.

The inspect import inside __init__ runs on every instantiation. Move it to the top of the file with the other imports.

♻️ Proposed refactor

+import inspect
+
 # ... other top-level imports ...

         # MONAI >= 1.3.0 backward compatibility handle
-        import inspect
         vit_params = inspect.signature(ViT.__init__).parameters

🤖 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 `@UNETR/BTCV/networks/unetr.py` at line 102, Move the inspect import from the
__init__ method to the module-level import section in unetr.py, alongside the
other imports, and remove the local import while leaving the existing inspect
usage unchanged.

🤖 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.

Inline comments:
In `@UNETR/BTCV/networks/unetr.py`:
- Line 37: Update the UNETR constructor around the proj_type parameter to accept
the legacy pos_embed keyword as an alias while retaining the new proj_type API.
Resolve both inputs consistently, with the explicitly provided compatibility
value taking effect when appropriate, and pass the resolved projection type
through the existing initialization path so callers in main.py and test.py
remain valid.

---

Outside diff comments:
In `@UNETR/BTCV/networks/unetr.py`:
- Around line 211-215: Update the bias checkpoint lookup in the weight-loading
block to use the same module path as the corresponding weight lookup, including
the patch_embedding segment before patch_embeddings.1.bias. Keep the target
module self.vit.patch_embedding.patch_embeddings[1] unchanged.

---

Nitpick comments:
In `@UNETR/BTCV/networks/unetr.py`:
- Line 102: Move the inspect import from the __init__ method to the module-level
import section in unetr.py, alongside the other imports, and remove the local
import while leaving the existing inspect usage unchanged.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Pro

Run ID: 987ebf72-a466-419c-b597-3a3cf445510a

📥 Commits

Reviewing files that changed from the base of the PR and between 21ed8e5 and a5ca63f.

📒 Files selected for processing (1)

• UNETR/BTCV/networks/unetr.py