torchcell.losses.multi_dim_nan_tolerant.MseCategoricalEntropyRegLoss

class MseCategoricalEntropyRegLoss(num_classes: int, num_tasks: int = 2, weights: Tensor | None = None, lambda_d: float = 0.1, lambda_t: float = 0.5)

Bases: Module

Combine MSE loss with entropy regularization, tracking per-dimension losses.

forward(continuous_pred: Tensor, continuous_target: Tensor, logits: Tensor, categorical_target: Tensor, pooled_features: Tensor) → tuple[Tensor, dict[str, Tensor]]

Forward pass tracking all loss components.

Returns total loss and dictionary of component losses.