torchcell.losses.multi_dim_nan_tolerant.NaNTolerantHuberLoss

class NaNTolerantHuberLoss(delta: float = 1.0)[source]

Bases: Module

Per-dimension Huber loss that ignores NaN targets.

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

Compute Huber loss while properly handling NaN values.

Parameters:

  • y_pred (torch.Tensor) – Predictions [batch_size, num_dims]

  • y_true (torch.Tensor) – Ground truth [batch_size, num_dims]

Returns:

(dim_means, mask) - Loss per dimension and validity mask

Return type:

tuple