From 02caedf7c586356ab0faf8785cd7d0d07a1fd675 Mon Sep 17 00:00:00 2001 From: Shizoqua Date: Mon, 6 Jul 2026 20:20:29 +0100 Subject: [PATCH] Fix NaN gradient in PerceptualLoss normalize_tensor Fixes #8412 Signed-off-by: Shizoqua --- monai/losses/perceptual.py | 6 ++++-- tests/losses/test_perceptual_loss.py | 11 +++++++++++ 2 files changed, 15 insertions(+), 2 deletions(-) diff --git a/monai/losses/perceptual.py b/monai/losses/perceptual.py index 635a3e75ce..fdd3d29b5c 100644 --- a/monai/losses/perceptual.py +++ b/monai/losses/perceptual.py @@ -311,8 +311,10 @@ def spatial_average_3d(x: torch.Tensor, keepdim: bool = True) -> torch.Tensor: def normalize_tensor(x: torch.Tensor, eps: float = 1e-10) -> torch.Tensor: - norm_factor = torch.sqrt(torch.sum(x**2, dim=1, keepdim=True)) - return x / (norm_factor + eps) + # Add eps inside the sqrt so the gradient stays finite when the norm is zero + # (e.g. identical input/target features), avoiding NaNs from SqrtBackward. See issue #8412. + norm_factor = torch.sqrt(torch.sum(x**2, dim=1, keepdim=True) + eps) + return x / norm_factor def medicalnet_intensity_normalisation(volume): diff --git a/tests/losses/test_perceptual_loss.py b/tests/losses/test_perceptual_loss.py index 8d94fdc1ae..051a79fef1 100644 --- a/tests/losses/test_perceptual_loss.py +++ b/tests/losses/test_perceptual_loss.py @@ -17,6 +17,7 @@ from parameterized import parameterized from monai.losses import PerceptualLoss +from monai.losses.perceptual import normalize_tensor from monai.utils import optional_import from tests.test_utils import assert_allclose, skip_if_downloading_fails, skip_if_quick @@ -126,6 +127,16 @@ def test_non_medicalnet_3d_without_fake_3d(self, network_type): with self.assertRaises(ValueError): PerceptualLoss(spatial_dims=3, network_type=network_type, is_fake_3d=False) + def test_normalize_tensor_zero_norm_finite_gradient(self): + # regression test for #8412: a zero-norm feature vector (e.g. from identical + # input/target features) must not produce NaN gradients via SqrtBackward. + x = torch.zeros(2, 4, 8, 8, requires_grad=True) + out = normalize_tensor(x) + out.sum().backward() + self.assertFalse(torch.isnan(out).any()) + self.assertIsNotNone(x.grad) + self.assertFalse(torch.isnan(x.grad).any()) + if __name__ == "__main__": unittest.main()