torchcell.losses.multi_dim_nan_tolerant.MultiDimNaNTolerantCELoss

class MultiDimNaNTolerantCELoss(num_classes: int, num_tasks: int = 2, eps: float = 1e-07)

Bases: Module

Multi-task cross-entropy loss that ignores NaN targets per task.

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

Compute Cross Entropy loss while handling NaN values.

Parameters:

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

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

Returns:

(dim_losses, mask)

Return type:

tuple