hyperboloid works after minor fix, but sphere is not working

[thoglu]

Hyperboloid: The hyperboloid code throws an error due to the "loc" property being defined after the parent init is called.
By putting https://github.com/CUAI/Neural-Manifold-Ordinary-Differential-Equations/blob/master/distributions/wnormal.py l.14 (the parent init call) to the end, we can solve that. After that the plots from the README can be constructed.

Sphere: The sphere does not seem to work at all. I am getting either black visualization or azimuthally symmetric ones.
Are there any crucial commits missing? It woul be nice to be able to reproduce the plots on the README, right now it is not possible.

Using pytorch 1.9 and the latest torchdiffeq.