torchcell.losses.multi_dim_nan_tolerant.NaNTolerantLogCoshLoss

class NaNTolerantLogCoshLoss[source]

Bases: Module

Per-dimension log-cosh loss that ignores NaN targets.

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

Compute log(cosh) loss while properly handling NaN values.

The log(cosh) loss is defined as: L = log(cosh(y_pred - y_true))

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