I've been experimenting with the Parametric UMAP using the example notebooks. However, I'm encountering an issue where the loss values I'm getting are significantly higher than in notebooks. This has occurred with both the second (02.0-parametric-umap-mnist-embedding-convnet.ipynb) and third (03.0-parametric-umap-mnist-embedding-convnet-with-reconstruction.ipynb) notebooks.