normdist2plane potential implementation error?

[yangshengaa]

Hi! It looks like when norm=True, normdist2plane is multiplied by the incorrect term?

In normdist2plane, there is a line
a_norm = a.norm(dim=dim, keepdim=keepdim, p=2).clamp_min(MIN_NORM)
this is the euclidean norm

and subsequently, when norm=True,
res = res * a_norm

However, res, I believe, should be multiplied by the Riemannian norm. From equation 24 in Hyperbolic Neural Network, the distance is computed by
$$sign(\langle -p_k \oplus_c x, a_k\rangle) \sqrt{g^c_{p_k}(a_k, a_k)} d_c(x, \tilde{H}_{a_k, p_k})$$

And $\sqrt{g^c_{p_k}(a_k, a_k)}$ shall take the place of a_norm in the multiplication.

Please correct me if I am wrong. I appreciate your help in advance.