ml.tasks.losses.diffusion

Defines loss functions used for Diffusion models.

ml.tasks.losses.diffusion.pseudo_huber_loss(x: Tensor, y: Tensor, dim: int = -1, factor: float = 0.00054, keepdim: bool = False) → Tensor[source]

Returns the pseudo-Huber loss.

This is taken from the Consistency Models paper.

Parameters:

  • x – The input tensor.

  • y – The target tensor.

  • dim – The dimension to compute the loss over.

  • factor – The factor to use in the loss.

  • keepdim – Whether to keep the dimension or not.

Returns:

The pseudo-Huber loss over the given dimension (i.e., that )