Unable to train custom model with ktrain. Model defined with custom layer and Functional API

Question

I have defined a model with custom layer to implement memory module. There is something wrong when ktrain save_weights of model. I think there is issue in saving we