swe_nonlinear_pino: physicsnemo_fdm_swe_nonlin slices batch dim ([1:-1]) -> empty PDE residual / NaN loss_pde with default batchsize=1 #2214

In examples/cfd/swe_nonlinear_pino/train_utils/losses.py, physicsnemo_fdm_swe_nonlin slices the PDE residuals along the batch dimension:

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

The reference implementation fdm_swe_nonlin in the same file trims the last dimension ([..., 1:-1]). The time trimming is also already applied to the inputs (1 : nt - 1, nt + 1 : 2 * nt - 1).

With the default batchsize: 1 in config_pino.yaml, the residual tensors become empty, e.g. Dh.shape == (0, 99, 128, 128):

• loss_pde is NaN;
• the PDE-term gradient is 0;
• training proceeds without any error, so the physics loss silently has no effect.

With batch size 2 the residuals are also empty. With batch size 3, only 1 of 3 samples is used.

Reproduce (v2.2.2, pip install nvidia-physicsnemo[sym], A100):

1. Run swe_nonlinear_pino with the default config and the physicsnemo derivative path.
2. The log shows test loss_pde = nan.

Fix: remove the slice (Dh = pde_Dh["Dh"], and the same for Du and Dv). After this change loss_pde is finite and decreases (train 0.206 → 0.023, test 0.019 after 2 epochs). It then matches the behaviour of loss.derivative=original.

The lines are unchanged on main as of 2026-11-01.