Thank you for the interest you have shown.
Actually, this example is a simple code that I wrote when I was learning FL for the first time. The aggregation method is valid for only one architecture. I will share my improvements on some functions and attack simulations under a different repository soon.