torchcell.losses.dango.PreThenPost

class PreThenPost(transition_epoch: int = 10)

Bases: DangoLossSched

Schedule that uses reconstruction loss only before transition_epoch, then switches to interaction loss only.

Parameters:

transition_epoch – Epoch at which to switch from reconstruction to interaction loss

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

Compute loss based on transition epoch.

Before transition: alpha = 1.0 (reconstruction only) After transition: alpha = 0.0 (interaction only)