torchcell.losses.multi_dim_nan_tolerant.CombinedCELoss

class CombinedCELoss(num_classes: int, num_tasks: int = 2, weights: Tensor | None = None)

Bases: Module

Cross-entropy loss combining per-task losses with dimension weights.

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

Compute weighted cross entropy loss.

Parameters:

• logits – Model predictions [batch_size, num_tasks * num_classes]

• targets – Target labels [batch_size, num_tasks * num_classes]

Returns:

(total_loss, dim_losses)

Return type:

tuple