fix: handle singleton channel dim in expand_as_one_hot (closes #36) - #38
Open
rtmalikian wants to merge 1 commit into
Open
fix: handle singleton channel dim in expand_as_one_hot (closes #36)#38rtmalikian wants to merge 1 commit into
rtmalikian wants to merge 1 commit into
Conversation
…017#36) When the target tensor arrives as 5D with a singleton channel dimension (e.g. [B, 1, D, H, W] from a DataLoader), expand_as_one_hot previously returned it unchanged because it only checked input.dim() == 5. This caused the assertion in _AbstractDiceLoss.forward() to fail since [B, C, D, H, W] != [B, 1, D, H, W]. Fix: when the input is 5D with 1 channel, squeeze to 4D before proceeding with one-hot expansion. If the input already has C channels, return as-is (preserving the original fast path). If the channel count is neither 1 nor C, raise a clear ValueError.
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
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
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.
Summary
Fixes #36 —
expand_as_one_hot()silently returns a 5D target tensor unchanged when it has a singleton channel dimension ([B, 1, D, H, W]), causing the shape assertion in_AbstractDiceLoss.forward()to fail.Root Cause
In
lib/losses3D/basic.py,expand_as_one_hot()has an early return:When the target arrives as
[B, 1, D, H, W](5D with a singleton channel from the DataLoader), this returns it unchanged. The subsequent assertioninput.size() == target.size()then fails because the model output is[B, C, D, H, W]but the target is[B, 1, D, H, W].Fix
Modified the 5D check to handle three cases:
C→ return as-is (already one-hot, original fast path)ValueErrorVerification
Tested with:
[B, D, H, W]— original behavior preserved ✅[B, 1, D, H, W]— correctly expanded to one-hot ✅[B, C, D, H, W]— returned as-is ✅input [4,4,128,128,48],target [4,1,128,128,48]— loss computed successfully ✅GeneralizedDiceLoss— also works with the fix ✅About the Author: Raphael Malikian — Clinical AI Solutions Architect. I specialise in building and fixing AI/ML systems for healthcare, including vector databases, RAG pipelines, and clinical NLP. If you need help with your project or think I can add value to your organisation, feel free to reach out — I'd love to connect.
📧 rtmalikian@gmail.com
🔗 GitHub: https://github.com/rtmalikian
🔗 LinkedIn: http://www.linkedin.com/in/raphael-t-malikian-mbbs-bsc-hons-71075436a
Disclosure: This code was developed with assistance from mimo-2.5-pro (Xiaomi) via Hermes Agent (Nous Research). All changes were reviewed, tested against the actual codebase, and verified for correctness.