Fix NaN gradient in PerceptualLoss normalize_tensor

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):

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()