Skip to content

fix: use register_buffer for kernel and kernel_vol in LocalNormalizedCrossCorrelationLoss #8818

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
Zeesejo wants to merge 1 commit into
base: dev
Choose a base branch
from
+2 −3
Open

fix: use register_buffer for kernel and kernel_vol in LocalNormalizedCrossCorrelationLoss #8818

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 2 additions & 3 deletions monai/losses/image_dissimilarity.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -111,9 +111,8 @@ def __init__(
raise ValueError(f"kernel_size must be odd, got {self.kernel_size}")

_kernel = look_up_option(kernel_type, kernel_dict)
self.kernel = _kernel(self.kernel_size)
self.kernel.require_grads = False
self.kernel_vol = self.get_kernel_vol()
self.register_buffer("kernel", _kernel(self.kernel_size))
self.register_buffer("kernel_vol", self.get_kernel_vol())
Comment on lines +114 to +115
Copy link
Copy Markdown
Contributor

@coderabbitai coderabbitai bot Apr 11, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟡 Minor

Add regression tests for the new buffer semantics.

Line 114 and Line 115 change state behavior, but current tests only validate numerics/error paths. Please add assertions that kernel and kernel_vol are present in state_dict, remain non-parameter buffers, and move with module device changes.

As per coding guidelines, "Ensure new or modified definitions will be covered by existing or new unit tests."

🤖 Prompt for AI Agents 
Verify each finding against the current code and only fix it if needed.

In `@monai/losses/image_dissimilarity.py` around lines 114 - 115, Add unit tests
that verify the new buffer semantics for the ImageDissimilarity module:
instantiate the module (using the class that calls self.register_buffer for
"kernel" and "kernel_vol"), assert both keys appear in module.state_dict(),
assert they are not in module.parameters() (i.e., non-parameter buffers), and
test device movement by .to(device) or .cuda() to confirm the tensors in
module.kernel and module.kernel_vol move to the target device; also include a
test that the values persist across state_dict save/load operations to ensure
correct buffer behavior. Ensure you reference the module/class name that defines
get_kernel_vol and register_buffer when adding the tests.

Comment on lines 113 to +115
Copy link

Copilot AI Apr 11, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Registering kernel/kernel_vol as persistent buffers changes this module’s state_dict() schema (older checkpoints without these keys will fail load_state_dict(strict=True) with missing keys). Consider whether backward compatibility is required; if so, handle missing keys (e.g., custom _load_from_state_dict) and/or make derived kernel_vol non-persistent and recompute it from kernel on load.

Copilot uses AI. Check for mistakes.
Comment on lines +114 to +115
Copy link

Copilot AI Apr 11, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This change is intended to improve device placement / serialization behavior, but there’s no regression test asserting (1) kernel/kernel_vol appear in state_dict() and/or (2) they follow .to(device/dtype) without per-forward casting. Adding a small unit test would help prevent this from regressing again.

Copilot uses AI. Check for mistakes.

self.smooth_nr = float(smooth_nr)
self.smooth_dr = float(smooth_dr)
Expand Down
Loading