Nan during GPU training

Dear Developers,

I'm encountering an issue during Parametric UMAP training on GPU: the loss becomes NaN, even when using the official demo notebooks without any modifications. My setup:

• GPU: NVIDIA GeForce RTX 4090
• CUDA/cuDNN: Detected and properly loaded
• TensorFlow version: I have tried everything from 2.16 to recent (currently 2.19 I think)
• I use the most recent version of umap-learn

The training works fine when I force the model to use the CPU. On GPU however, the training starts normally, but quickly diverges and logs loss: nan by the first epoch. I have tried ta few things based on other issues and discussions:

• Lowering the learning rate (e.g., 1e-3, 1e-4)
• Disabling XLA (tf.config.optimizer.set_jit(False))
• Disabling mixed precision
• Using .astype('float64') for fit_transform
• Reducing the batch size

Despite all of the above, the issue persists on GPU. Any guidance on resolving this would be appreciated.

Best regards