torchcell.losses.multi_dim_nan_tolerant.NaNTolerantMSELoss

class NaNTolerantMSELoss

Bases: Module

Per-dimension MSE loss that ignores NaN targets.

forward(y_pred: Tensor, y_true: Tensor) → tuple[Tensor, Tensor]

Compute MSE loss while properly handling NaN values.