Load weights of keras model with custom layers

Question

I am trying to save the model and weights for the WGAN implementation found in

https://github.com/keras-team/keras-contrib/blob/master/examples/improved_wgan.py .

However, when we perform load_weights(...), it throws "axes don't match array" error. I tried to use the save model, load model with custom objects. But that throws the error too. Can anyone please help me to save the weights properly when we have a custom loss and custom layers in Keras?

Answer 1

If someone needs the solution, saving the model in yaml format and then the weights in h5 worked for me!