How to compute the mean of weights of multiple models?

Hi i'm a student and i'm working on a Federated Learning problem, but before doing that with the proper tools like OpenFL or Flower, I started a little experiment to try in local to train using this technique.

I managed to train multiple models using IID data, now I'm struggling with the local_update() function that should collect the models and then i need to take all the weights of these models and compute their mean. I read some documentation of Keras and Tensorflow that I'm using for my work, and i found some functions but i can't get it to work properly.

Currently this is my local_update() that's not working

    def local_update(self, models):

        
        weights = []
        #Take the weights of the models and compute the mean then return the weights to an updated model

        for model in models:
            for layer in model.layers:
                weights = layer.get_weights()
        
        #Compute the mean of weights  
        weights = np.mean(weights, axis=0)


        for layer in self.model.layers:
            self.model.set_weights(weights)     


        return self.model

Thank you in advance for the help.