[dmitriySem]

closes #2214

Problem

In examples/cfd/swe_nonlinear_pino/train_utils/losses.py, the function physicsnemo_fdm_swe_nonlin slices the PDE residual tensors Dh, Du, and Dv along the batch dimension using [1:-1]. With the default batchsize: 1 in the example config, this produces an empty tensor (shape (0, ...)), causing loss_pde to become NaN and the physics loss gradient to be zero. Training proceeds silently without any physics contribution.

This is inconsistent with the reference function fdm_swe_nonlin in the same file, which correctly slices along the last (spatial) dimension using [..., 1:-1].

Fix

Removed the batch-dimension slicing from physicsnemo_fdm_swe_nonlin:

- Dh = pde_Dh["Dh"][1:-1]
- Du = pde_Du["Du"][1:-1]
- Dv = pde_Dv["Dv"][1:-1]
+ Dh = pde_Dh["Dh"]
+ Du = pde_Du["Du"]
+ Dv = pde_Dv["Dv"]

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[github-actions]

CODEOWNERS review map

Current for commit 0a352e31260c. An approval covers every file listed for that owner; one owner is sufficient for shared files.

No CODEOWNER

• CHANGELOG.md
• examples/cfd/swe_nonlinear_pino/train_utils/losses.py

Comment /codeowners-info to refresh.

[greptile-apps]

[Medium impact] The change affects the shallow-water PINO example's PhysicsNeMo training path.

The fix appears safe to merge, with a non-blocking request for regression coverage.

Findings

1.  Batch fix lacks regression coverage ▶

Summary

Fixes the shallow-water physics loss by keeping every batch element in Dh, Du, and Dv.

• Shallow-water training keeps PDE residuals from one-item batches.

_{Reviews (1) · Last reviewed commit: "Fixed batch-dimension slicing in `physic..." · Reviewed by Greptile}

[greptile-apps]

Batch fix lacks regression coverage

The batch-slicing fix has no regression test. Add a small test for physicsnemo_fdm_swe_nonlin and physicsnemo_fourier with batch sizes 1, 2, and 3. Check that each residual keeps every batch element and that a nonconstant input produces a finite loss and finite, nonzero gradients. This would catch both empty residuals for small batches and silently dropped samples for larger batches.

Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!

[dmitriySem]

Thanks for the concrete suggestion. I agree that a regression test would be valuable. Since this is my first contribution to PhysicsNeMo and the current PR is intentionally minimal, I'd like to get feedback from the project maintainers first. If a test is required for acceptance, I'll add it to this PR or prepare a separate one.

[greptile-apps]

That makes sense. My comment was a coverage suggestion, not a claim that a regression test is required for acceptance. Keeping this PR focused and waiting for maintainer feedback is reasonable; if they request coverage, you can add it here or in a follow-up.

Tip: You can customize Greptile's behavior for this repo with .greptile/rules.md and .greptile/config.json.