torchcell.losses.diffusion_loss.DiffusionLoss

class DiffusionLoss(model: Module, *, lambda_diffusion: float = 1.0, lambda_x0: float = 0.0, x0_loss: Module | None = None)

Bases: Module

Strict diffusion loss for conditional decoders.

What it does:

• computes the model’s diffusion loss

• (optionally) adds explicit x0 supervision at sampled timesteps

No auxiliary terms, no silent fallbacks.

forward(predictions: Tensor, targets: Tensor, context: Tensor | None = None, *, epoch: int | None = None, t_mode: str = 'full') → tuple[Tensor, dict[str, Tensor]]

Compute the weighted diffusion (and optional x0) loss and a breakdown.