torchcell.losses.dango.DangoLossSched

class DangoLossSched[source]

Bases: ABC

Abstract base class for DANGO loss schedulers.

All schedulers must implement a forward method that computes the total loss based on reconstruction loss, interaction loss, and current epoch.

abstractmethod forward(recon_loss: Tensor, interaction_loss: Tensor, current_epoch: int = 0) → tuple[Tensor, float, Tensor, Tensor][source]

Compute the combined loss according to the schedule.

Parameters:

  • recon_loss – Reconstruction loss value

  • interaction_loss – Interaction prediction loss value

  • current_epoch – Current training epoch

Returns:

  • total_loss: The combined loss value

  • alpha: Weight of reconstruction loss (for logging)

  • weighted_recon_loss: The weighted reconstruction loss component

  • weighted_interaction_loss: The weighted interaction loss component

Return type:

Tuple containing