ONNX export bug

[brunoalvarez5]

My comment is on the file parametric_umap.py

I've found an error when saving the model into an ONNX file, since it adds a ReLu at the end of the model. I've found the reason for this is that when defining the original encoder in the function forward it does not have this, however when it has to save it into ONNX first it passes the model to pytorch. So when it checks if the encoder is None and then defines it, in this definition it adds a ReLu and then exports the model like that.
I have attached screenshots of the two sections of the code I'm referring to for better clarity. (My solution was only commenting the line where the ReLu is defined) I really like this project, so could I perhaps open a fork to help with the project and fix this bug?
Thank you.