Fix NaN gradient in PerceptualLoss normalize_tensor

[Shizoqua]

Fixes #8412

Description

PerceptualLoss could crash with RuntimeError: Function 'SqrtBackward0' returned nan values in its 0th output when the loss became very low (near-identical input/target features). The cause was in normalize_tensor, where eps was added after the square root — this guarded the forward division but not the sqrt gradient, so a zero feature norm produced 1/(2·√0) = inf → NaN during backprop. This PR moves eps inside the sqrt (sqrt(sum(x**2) + eps)), keeping the forward output unchanged for normal inputs while ensuring the gradient stays finite. This matches the standard LPIPS normalization pattern.

Types of changes

• Non-breaking change (fix or new feature that would not break existing functionality).
• Breaking change (fix or new feature that would cause existing functionality to change).
• New tests added to cover the changes.
• Integration tests passed locally by running ./runtests.sh -f -u --net --coverage.
• Quick tests passed locally by running ./runtests.sh --quick --unittests --disttests.
• In-line docstrings updated.
• Documentation updated, tested make html command in the docs/ folder.

[coderabbitai]

📝 Walkthrough

Walkthrough

The change modifies normalize_tensor in monai/losses/perceptual.py to add eps inside the square root when computing the norm factor, rather than adding eps to the denominator afterward, preventing NaN gradients when the norm is zero (e.g., identical input/target features). A corresponding regression test is added to tests/losses/test_perceptual_loss.py verifying finite output and gradients for zero-norm tensors.

Estimated code review effort: 2 (Simple) | ~10 minutes

Changes

File	Summary
monai/losses/perceptual.py	Moved eps inside sqrt in normalize_tensor for numerical stability.
tests/losses/test_perceptual_loss.py	Added import and regression test for zero-norm gradient stability.

Related issues: Fixes PerceptualLoss NaN gradient bug (#8412) caused by SqrtBackward0 returning NaN when loss is near zero.

Suggested reviewers: ericspod, KumoLiu

Poem

A rabbit hopped through gradient trees,
Found NaNs hiding in the sqrt's unease.
Moved the eps, tucked it in tight,
Now zero norms compute just right.
Thump thump — tests pass, all is bright.

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title is concise and accurately summarizes the main fix: preventing NaN gradients in PerceptualLoss normalization.
Description check	✅ Passed	The description follows the template and includes the issue reference, change summary, and type-of-changes checklist.
Linked Issues check	✅ Passed	The code and test changes directly address issue #8412 by moving eps inside the sqrt and verifying gradients stay finite.
Out of Scope Changes check	✅ Passed	The PR stays focused on the reported normalization bug and its regression test, with no unrelated changes apparent.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}

[coderabbitai]

🧹 Nitpick comments (1)

monai/losses/perceptual.py (1)

313-317: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Correct fix for NaN gradient; add missing docstring.

Moving eps inside the sqrt is the right fix—derivative of sqrt at 0 is undefined/infinite, so keeping the argument bounded away from 0 avoids the NaN. Comments explain rationale well.

However, normalize_tensor still lacks a docstring describing x, eps, and the return value.

📝 Suggested docstring

 def normalize_tensor(x: torch.Tensor, eps: float = 1e-10) -> torch.Tensor:
+    """Normalize a tensor across the channel dimension (dim=1).
+
+    Args:
+        x: input tensor to normalize.
+        eps: small value added inside the sqrt to avoid NaN gradients when the
+            per-sample norm is zero (see issue `#8412`).
+
+    Returns:
+        The channel-normalized tensor with the same shape as ``x``.
+    """
     # 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

As per path instructions, "Docstrings should be present for all definition which describe each variable, return value, and raised exception in the appropriate section of the Google-style of docstrings."

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@monai/losses/perceptual.py` around lines 313 - 317, The normalize_tensor
function in perceptual.py needs a Google-style docstring. Add a concise
docstring to normalize_tensor that describes the x tensor argument, the eps
parameter, and the normalized tensor return value, matching the project’s
documentation standard while keeping the existing NaN-safe implementation
intact.

Source: Path instructions

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Nitpick comments:
In `@monai/losses/perceptual.py`:
- Around line 313-317: The normalize_tensor function in perceptual.py needs a
Google-style docstring. Add a concise docstring to normalize_tensor that
describes the x tensor argument, the eps parameter, and the normalized tensor
return value, matching the project’s documentation standard while keeping the
existing NaN-safe implementation intact.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

Run ID: 28a5a6b9-7313-4f8f-ac96-8d3344078194

📥 Commits

Reviewing files that changed from the base of the PR and between 229f519 and 02caedf.

📒 Files selected for processing (2)

• monai/losses/perceptual.py
• tests/losses/test_perceptual_loss.py