Train Keras model after Split for federated learning

Question

Hi i am trying to train the keras model after split. simulating the keras model are placed on different devices and each device have their weights. that devices passes their intermediate output to host and host device receive intermediate output and send back gradient.How can we train the splited model here is the example of splited model for mnist dataset

model1 = keras.Sequential(
[
    keras.Input(shape=input_shape),
    layers.Conv2D(32, kernel_size=2, activation="relu",strides=1, padding="same"),
    layers.MaxPooling2D(pool_size=2),
]
)
model2 = keras.Sequential(
[
    
    layers.Conv2D(64, kernel_size=2, activation="relu", strides=1, padding="same"),
    layers.MaxPooling2D(pool_size=2),
    layers.Flatten(),
    layers.Dropout(0.5),
    layers.Dense(num_classes, activation="softmax"),
]
)

I am expecting the way to train these models on different devices where guest device will have images and model1 and send the intermediate output to host after predict with moodel1. Host contain model2 and label for images. host 2 will receive intermediate output and host will train model2 and send the gradient to guest which will train the model1.