Fix hardcoded input dim in DiffusionModelEncoder #8514
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
WalkthroughThe initialization of the output layer in the Changes
Sequence Diagram(s)sequenceDiagram
participant User
participant DiffusionModelEncoder
participant nn.Sequential
User->>DiffusionModelEncoder: forward(input)
DiffusionModelEncoder->>DiffusionModelEncoder: Flatten features (h)
alt self.out not initialized
DiffusionModelEncoder->>nn.Sequential: Create with input size = h.shape[1]
DiffusionModelEncoder->>DiffusionModelEncoder: Assign to self.out
end
DiffusionModelEncoder->>nn.Sequential: Pass h through self.out
nn.Sequential-->>DiffusionModelEncoder: Output tensor
DiffusionModelEncoder-->>User: Return output
Poem
📜 Recent review detailsConfiguration used: CodeRabbit UI 📒 Files selected for processing (1)
🚧 Files skipped from review as they are similar to previous changes (1)
✨ Finishing Touches
🪧 TipsChatThere are 3 ways to chat with CodeRabbit:
SupportNeed help? Create a ticket on our support page for assistance with any issues or questions. Note: Be mindful of the bot's finite context window. It's strongly recommended to break down tasks such as reading entire modules into smaller chunks. For a focused discussion, use review comments to chat about specific files and their changes, instead of using the PR comments. CodeRabbit Commands (Invoked using PR comments)
Other keywords and placeholders
CodeRabbit Configuration File (
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Nitpick comments (1)
monai/networks/nets/diffusion_model_unet.py (1)
2053-2056: Consider device placement and thread safety for lazy initialization.The lazy initialization logic is correct, but consider these improvements:
- Device placement: The dynamically created module should be moved to the same device as the input tensor.
- Thread safety: In multi-threaded environments, this could create race conditions.
Apply this diff to ensure proper device placement:
- if self.out is None: - self.out = nn.Sequential( - nn.Linear(h.shape[1], 512), nn.ReLU(), nn.Dropout(0.1), nn.Linear(512, self.out_channels) - ) + if self.out is None: + self.out = nn.Sequential( + nn.Linear(h.shape[1], 512), nn.ReLU(), nn.Dropout(0.1), nn.Linear(512, self.out_channels) + ).to(h.device)For thread safety, consider using a lock or moving the initialization to a separate method called during model setup.
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (1)
monai/networks/nets/diffusion_model_unet.py(2 hunks)
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (18)
- GitHub Check: min-dep-pytorch (2.6.0)
- GitHub Check: min-dep-py3 (3.12)
- GitHub Check: min-dep-pytorch (2.5.1)
- GitHub Check: min-dep-py3 (3.11)
- GitHub Check: min-dep-pytorch (2.4.1)
- GitHub Check: min-dep-os (ubuntu-latest)
- GitHub Check: min-dep-py3 (3.10)
- GitHub Check: min-dep-py3 (3.9)
- GitHub Check: min-dep-os (windows-latest)
- GitHub Check: min-dep-os (macOS-latest)
- GitHub Check: quick-py3 (macOS-latest)
- GitHub Check: build-docs
- GitHub Check: flake8-py3 (mypy)
- GitHub Check: quick-py3 (ubuntu-latest)
- GitHub Check: quick-py3 (windows-latest)
- GitHub Check: packaging
- GitHub Check: flake8-py3 (codeformat)
- GitHub Check: flake8-py3 (pytype)
🔇 Additional comments (1)
monai/networks/nets/diffusion_model_unet.py (1)
2008-2008: LGTM: Proper lazy initialization setup.Setting
self.out = Nonecorrectly defers the output layer creation until runtime, enabling dynamic input dimension adaptation.
Signed-off-by: IamTingTing <[email protected]>
IamTingTing
force-pushed
the
fix-encoder-dim
branch
from
3e45145 to
595ca84
Compare
virginiafdez
approved these changes
Jul 22, 2025
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Fixes #8496
Description
A few sentences describing the changes proposed in this pull request.
Types of changes
./runtests.sh -f -u --net --coverage../runtests.sh --quick --unittests --disttests.make htmlcommand in thedocs/folder.Summary by CodeRabbit