torchcell.losses.isomorphic_cell_loss.ICLossStd

class ICLossStd(lambda_dist: float, lambda_supcr: float, lambda_reg: float = 0.01, init_sigma: float = 1.0, task_weights: Tensor | None = None, eps: float = 1e-06)[source]

Bases: Module

IC loss with learnable per-task standard deviations (heteroscedastic).

forward(predictions: Tensor, targets: Tensor, z_P: Tensor, z_I: Tensor) → tuple[Tensor, dict[str, Any]][source]

Forward pass to compute the heteroscedastic IC loss.

Parameters:

  • predictions – Model predictions [batch_size, 2]

  • targets – Target values [batch_size, 2]

  • z_P – Perturbed difference embeddings from the model

  • z_I – Intact embeddings from the model

Returns:

(total_loss, loss_dict)

Return type:

tuple