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.