torchcell.losses.mle_wasserstein.BufferedWeightedSupCRCell

class BufferedWeightedSupCRCell(buffer_size: int = 256, embedding_dim: int = 128, temperature: float = 0.1, weights: Tensor | None = None, min_samples: int = 32)

Bases: Module

Enhanced SupCR loss with circular buffer for small batch training.

update_buffer(embeddings: Tensor, labels: Tensor) → None

Update circular buffer with new samples.

get_buffer_samples() → tuple[Tensor, Tensor]

Get all valid samples from buffer.

forward(embeddings: Tensor, labels: Tensor, all_embeddings: Tensor | None = None, all_labels: Tensor | None = None, buffer_weight: float = 1.0, temperature: float | None = None) → tuple[Tensor, Tensor]

Compute SupCR loss with buffer support.

Parameters:

• embeddings – Local batch embeddings

• labels – Local batch labels

• all_embeddings – Gathered embeddings from all GPUs (optional)

• all_labels – Gathered labels from all GPUs (optional)

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

• temperature – Optional temperature override