Softmax tilts to one side (trained label) after aggregator mean computation

[graykode]

Hello Authors! Thanks for sharing nice work. Your code are awesome. I try your emoji_hunt demo with my trained model. I think this is important issue to implement Ferderated Learning. After mean computation, Softmax tilts to one side (trained label ) was occured

I guess this is ocurring when mean aggreation computation.