Class MappingHuber

flowmap/loss/mapping/mapping_huber.py:18–34

class MappingHuber(Mapping[MappingHuberCfg]):
    def forward_undistorted(
        self,
        delta: Float[Tensor, "*batch 2"],
    ) -> Float[Tensor, " *batch"]:
        norm = delta.norm(dim=-1)

        mapped = F.huber_loss(
            norm,
            torch.zeros_like(norm),
            reduction="none",
            delta=self.cfg.delta,
        )

        # Divide by the delta so that the gradient magnitude in the linear region
        # matches that of a regular L1 loss.
        return mapped / self.cfg.delta