torchcell.losses.mle_wasserstein.BufferedWeightedWassersteinLoss

class BufferedWeightedWassersteinLoss(buffer_size: int = 256, blur: float = 0.05, p: float = 2, scaling: float = 0.9, weights: Tensor | None = None, min_samples: int = 64)

Bases: Module

Enhanced Wasserstein loss with circular buffer for small batch training.

update_buffer(predictions: Tensor, targets: Tensor) → None

Update circular buffer with new samples.

get_buffer_samples() → tuple[Tensor, Tensor]

Get all valid samples from buffer.

forward(predictions: Tensor, targets: Tensor, all_predictions: Tensor | None = None, all_targets: Tensor | None = None, buffer_weight: float = 1.0) → tuple[Tensor, Tensor]

Compute Wasserstein loss with buffer support.

Parameters:

• predictions – Local batch predictions

• targets – Local batch targets

• all_predictions – Gathered predictions from all GPUs (optional)

• all_targets – Gathered targets from all GPUs (optional)

• buffer_weight – Weight for buffer contribution (0-1)