def diffusion_loss(dvf):
    """
    Calculate diffusion loss as a regularisation on the displacement vector field (DVF)

    Args:
        dvf: (Tensor of shape (N, 2, H, W)) displacement vector field estimated

    Returns:
        diffusion_loss_2d: (Scalar) diffusion regularisation loss
        """

    # spatial derivatives
    dvf_dx = dvf[:, :, 1:, 1:] - dvf[:, :, :-1, 1:]  # (N, 2, H-1, W-1)
    dvf_dy = dvf[:, :, 1:, 1:] - dvf[:, :, 1:, :-1]  # (N, 2, H-1, W-1)
    return (dvf_dx.pow(2) + dvf_dy.pow(2)).mean()