About switch_gate

[Heihaierr]

Hi, I'm trying to implement a simpler version of switch transformer following your work. But the detail of switch_gate is invisible, like limit_by_capacity. My implementation has a slight different result with fastmoe.

Can you release the detail code of switch_gate?

Thanks.

[laekov]

The pruning function is implemented here.