Get ValueError when apply transfer learning in federated learning (TFF)

I want to use pre_trained model in federated learning as following code:

first I build my model and set the weights on model and then I freeze convolutional layers and remove 4 last layer.

def create_keras_model():
    model = Sequential()
    model.add(Conv2D(16, kernel_size=(3,3), strides=(1,1), padding='same', activation='relu', input_shape=(226,232,1)))
    model.add(MaxPooling2D((2,2), strides=(2,2), padding='same'))
    model.add(Conv2D(32, kernel_size=(3,3), strides=(1,1), padding='same', activation='relu'))
    model.add(MaxPooling2D((2,2), strides=(2,2), padding='same'))
    model.add(Conv2D(64, kernel_size=(3,3), strides=(1,1), padding='same', activation='relu'))
    model.add(MaxPooling2D((2,2), strides=(2,2), padding='same'))
    model.add(Conv2D(64, kernel_size=(3,3), strides=(1,1), padding='same', activation='relu'))
    model.add(MaxPooling2D((2,2), strides=(2,2), padding='same'))
    model.add(Conv2D(128, kernel_size=(3,3), strides=(1,1), padding='same', activation='relu'))
    model.add(MaxPooling2D((2,2), strides=(2,2), padding='same'))
    model.add(Conv2D(128, kernel_size=(3,3), strides=(1,1), padding='same', activation='relu'))
    model.add(MaxPooling2D((2,2), strides=(2,2), padding='same'))
    model.add(Conv2D(256, kernel_size=(3,3), strides=(1,1), padding='same', activation='relu'))
    model.add(MaxPooling2D((2,2), strides=(2,2), padding='same'))
    model.add(Flatten())
    model.add(Dense(256, activation='relu'))
    model.add(Dense(128, activation='relu'))
    model.add(Dense(64, activation='relu'))
    model.add(Dense(10, activation='softmax'))
    return model

keras_model = create_keras_model()

server_state=FileCheckpointManager(root_dir= '/content/drive/MyDrive',
    prefix=  'federated_clustering',
    step= 1,
    keep_total= 1,
    keep_first= True).load_checkpoint(structure=server_state,round_num=10)

keras_model.set_weights(server_state)

for layer in keras_model.layers[:-4]:
  layer.trainable = False

model_pre = Model(inputs=keras_model.input,outputs=keras_model.layers[14].output)

next, I build new model.

def create_keras_model1():
    model = Sequential()
    model.add(model_pre)
    model.add(Dense(256, activation='relu'))
    model.add(Dense(128, activation='relu'))
    model.add(Dense(64, activation='relu'))
    model.add(Dense(2, activation='softmax'))
    return model

def model_fn():
  # We _must_ create a new model here, and _not_ capture it from an external
  # scope. TFF will call this within different graph contexts.
  keras_model = create_keras_model1()
  return tff.learning.from_keras_model(
      keras_model,
      input_spec=preprocessed_example_dataset.element_spec,
      loss=tf.keras.losses.SparseCategoricalCrossentropy(),
      metrics=[tf.keras.metrics.SparseCategoricalAccuracy()])

but I get ValueError when I want to use tff.learning.build_federated_averaging_process.

iterative_process = tff.learning.build_federated_averaging_process(
    model_fn,
    client_optimizer_fn=lambda: tf.keras.optimizers.SGD(learning_rate=0.001),
    server_optimizer_fn=lambda: tf.keras.optimizers.SGD(learning_rate=1.0))

ValueError: Your Layer or Model is in an invalid state. This can happen for the following cases:
 1. You might be interleaving estimator/non-estimator models or interleaving models/layers made in tf.compat.v1.Graph.as_default() with models/layers created outside of it. Converting a model to an estimator (via model_to_estimator) invalidates all models/layers made before the conversion (even if they were not the model converted to an estimator). Similarly, making a layer or a model inside a a tf.compat.v1.Graph invalidates all layers/models you previously made outside of the graph.
2. You might be using a custom keras layer implementation with  custom __init__ which didn't call super().__init__.  Please check the implementation of <class 'keras.engine.functional.Functional'> and its bases.

please help me to fix it.