torchcell.losses.dango.LinearUntilUniform

class LinearUntilUniform(transition_epoch: int = 20)

Bases: DangoLossSched

Schedule that linearly transitions from reconstruction-focused to uniform weighting.

Parameters:

transition_epoch – Epoch at which weights become uniform (0.5 each)

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

Compute loss with linearly decreasing alpha.

Alpha starts at 1.0 (reconstruction only) and linearly decreases to 0.5 (uniform) over transition_epoch epochs. After that, it remains at 0.5.