torchcell.losses.mle_wasserstein.WeightedWassersteinLoss

class WeightedWassersteinLoss(blur: float = 0.05, p: float = 2, scaling: float = 0.9, weights: Tensor | None = None)

Bases: Module

Weighted Wasserstein distance loss for multi-dimensional outputs.

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

Compute weighted Wasserstein distance.

Parameters:

• predictions – Predicted values [batch_size, num_dims]

• targets – Target values [batch_size, num_dims]

Returns:

Tuple of (total_loss, per_dimension_losses)