torchcell.losses.multi_dim_nan_tolerant.WeightedMSELoss

class WeightedMSELoss(weights: Tensor | None = None)

Bases: Module

Per-dimension MSE loss with optional normalized dimension weights.

forward(predictions: Tensor, targets: Tensor) → tuple[Tensor, Tensor]

Return the weighted MSE and per-dimension losses, ignoring NaN targets.