torchcell.losses.multi_dim_nan_tolerant.WeightedSupCRCell

class WeightedSupCRCell(temperature: float = 0.1, weights: Tensor | None = None, eps: float = 1e-07)

Bases: Module

SupCR loss with per-dimension weights for cell-state embeddings.

forward(perturbed_embeddings: Tensor, labels: Tensor) → tuple[Tensor, Tensor]

Return the weighted SupCR loss and per-dimension losses, ignoring NaNs.