From 6df895b2cdd68f095a3900ffb0358fbcee921040 Mon Sep 17 00:00:00 2001 From: Zeeshan Modi <92383127+Zeesejo@users.noreply.github.com> Date: Sat, 11 Apr 2026 19:37:44 +0200 Subject: [PATCH] fix: use register_buffer for kernel and kernel_vol in LocalNormalizedCrossCorrelationLoss Refactor kernel initialization to use register_buffer for better state management. Signed-off-by: Zeeshan Modi <92383127+Zeesejo@users.noreply.github.com> --- monai/losses/image_dissimilarity.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/monai/losses/image_dissimilarity.py b/monai/losses/image_dissimilarity.py index 29cce74926..3b68fead44 100644 --- a/monai/losses/image_dissimilarity.py +++ b/monai/losses/image_dissimilarity.py @@ -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()) self.smooth_nr = float(smooth_nr) self.smooth_dr = float(smooth_dr)