Failed to save a trained Parametric UMAP model ()

[mohammadm1985]

Hello,

I am trying to save a trained model (trained with fit.transform() function) with the built-in save function, but received this error:

WARNING:tensorflow:Compiled the loaded model, but the compiled metrics have yet to be built. `model.compile_metrics` will be empty until you train or evaluate the model.
WARNING:tensorflow:Compiled the loaded model, but the compiled metrics have yet to be built. `model.compile_metrics` will be empty until you train or evaluate the model.
INFO:tensorflow:Assets written to: /home/mmotall/complaints_subcat/model_training/models_developped/bert/model-save-test/par_umap2/encoder/assets
INFO:tensorflow:Assets written to: /home/mmotall/complaints_subcat/model_training/models_developped/bert/model-save-test/par_umap2/encoder/assets
Keras encoder model saved to /home/mmotall/complaints_subcat/model_training/models_developped/bert/model-save-test/par_umap2/encoder
WARNING:absl:Found untraced functions such as _update_step_xla while saving (showing 1 of 1). These functions will not be directly callable after loading.
INFO:tensorflow:Assets written to: /home/mmotall/complaints_subcat/model_training/models_developped/bert/model-save-test/par_umap2/parametric_model/assets
INFO:tensorflow:Assets written to: /home/mmotall/complaints_subcat/model_training/models_developped/bert/model-save-test/par_umap2/parametric_model/assets
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
/tmp/ipykernel_102/1354878626.py in <module>
      3 import keras
      4 
----> 5 embedder.save("/home/mmotall/complaints_subcat/model_training/models_developped/bert/model-save-test/par_umap2")

~/venv/lib/python3.7/site-packages/umap/parametric_umap.py in save(self, save_location, verbose)
    506         if self.parametric_model is not None:
    507             parametric_model_output = os.path.join(save_location, "parametric_model")
--> 508             self.parametric_model.save(parametric_model_output)
    509             if verbose:
    510                 print("Keras full model saved to {}".format(parametric_model_output))

~/venv/lib/python3.7/site-packages/keras/utils/traceback_utils.py in error_handler(*args, **kwargs)
     68             # To get the full stack trace, call:
     69             # `tf.debugging.disable_traceback_filtering()`
---> 70             raise e.with_traceback(filtered_tb) from None
     71         finally:
     72             del filtered_tb

~/venv/lib/python3.7/site-packages/keras/saving/experimental/serialization_lib.py in _get_class_or_fn_config(obj)
    142     else:
    143         raise TypeError(
--> 144             f"Cannot serialize object {obj} of type {type(obj)}. "
    145             "To be serializable, "
    146             "a class must implement the `get_config()` method."

TypeError: Cannot serialize object <tensorflow.python.eager.polymorphic_function.polymorphic_function.Function object at 0x7fbc2710f390> of type <class 'tensorflow.python.eager.polymorphic_function.polymorphic_function.Function'>. To be serializable, a class must implement the `get_config()` method.

I tried the sample given in parametric_umap official webpage and got the same error. This is crucial for me to find a solution and save the model as this is just a small part of a bigger model and the entire model will be useless if I cannot save and load the model. Would you please let me know if there is a way to make it work?

Update: I noticed that I can save a parametric umap model after training, but I cannot save it again or save any other instances of parametric umap model!!! When I try to save any other instances I get the above error. Meaning that I need to restart the kernel and run another model and save it. Do you have any idea?