Example of serialization error in keras model

minimal_example.py

seed = 10
import random
random.seed(seed)
import numpy as np
np.random.seed(seed)
import tensorflow as tf
tf.reset_default_graph()
tf.random.set_random_seed(seed)

import tensorflow.keras.backend as K
import tensorflow.keras.layers as kl
import tensorflow.keras.activations as ka
import tensorflow.keras.losses as klo
import tensorflow.keras.models as km
import tensorflow.keras.optimizers as ko

def create_model(X_dim):
    X_input = kl.Input(shape=(X_dim,))
    X = X_input
    X = kl.Dense(10)(X)
    X = kl.Dense(1)(X)
    X = K.expand_dims(X, axis=2)
    X = K.squeeze(X, axis=-1)
    y_output = X
    model = km.Model(inputs=X_input, outputs=y_output)
    optimizer = ko.Adam(learning_rate=0.01)
    loss = klo.MeanSquaredError()
    model.compile(optimizer, loss)
    return model

X_dim = 3
X = np.random.randn(100, X_dim)
y = np.random.randn(100)
model = create_model(X_dim)

yp = model.predict(X)
model.fit(X, y, verbose=10)
model.save('/tmp/here')

So, did you have a good solution?

My code is

parallel_model = keras.utils.multi_gpu_model(model, gpus=2)
parallel_model.compile(tf.optimizers.Adam(), loss='categorical_crossentropy', metrics=['accuracy'])
parallel_model.summary()
save_path = "./results/model.h5"
tb_callback = keras.callbacks.TensorBoard(args.log_dir, histogram_freq=0.1, write_graph=True,
                                          write_grads=True, write_images=True, embeddings_freq=0.5, update_freq='batch')
history = parallel_model.fit(x=x_train, y=y_train, batch_size=args.batch_size, epochs=args.epochs,
                             callbacks=[tb_callback], validation_split=args.fraction_validation, shuffle=True)
model.save(save_path)