Use Cnn matching as a black box

[chrisdok43]

I want to implement FedMA in a FML framework, is there a function inside this repo that i can use as a black box?
I want it to have input lets say two layers from client cnn and one from the global and to return the matched output.