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]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro Plus

Run ID: b7c00d65-4f11-41e1-bd57-cb5ca6a8fdbf

📥 Commits

Reviewing files that changed from the base of the PR and between 02caedf and 5304e6e.

📒 Files selected for processing (1)

• monai/losses/perceptual.py

🚧 Files skipped from review as they are similar to previous changes (1)

• monai/losses/perceptual.py

---

📝 Walkthrough

Walkthrough

The change updates normalize_tensor to add eps inside the square root used for norm computation. It removes the previous denominator adjustment. A regression test verifies finite outputs and gradients for a zero-valued tensor after backpropagation.

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

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly identifies the NaN gradient fix in PerceptualLoss normalization.
Description check	✅ Passed	The description explains the cause, solution, impact, linked issue, and regression test coverage.
Linked Issues check	✅ Passed	The code and regression test directly address issue #8412 by preventing NaN gradients at zero feature norms.
Out of Scope Changes check	✅ Passed	The changes are limited to the requested normalization fix and its focused regression test.
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

[vikashg]

Clean fix, correct math. Moving eps inside sqrt() ensures d/dx sqrt(x² + eps) is always finite — standard numerical stability pattern for normalized features. Regression test covers the zero-norm case. All CI passing. LGTM.

[vikashg]

Clean fix, correct math. Moving eps inside sqrt() ensures d/dx sqrt(x² + eps) is always finite — standard numerical stability pattern for normalized features. Regression test covers the zero-norm case. All CI passing. LGTM.