[BUG] SwitchGate in Switch MoE runtime error when computing the aux loss

[Mark-Sky]

Describe the bug
I run the switch_moe.py in the example fold. When I set the use_aux_loss=True, runtime error occurs:
RuntimeError: The size of tensor a (128) must match the size of tensor b (32) at non-singleton dimension 0

It is because the code in the SwitchGate class:

The load shape is different from the importance shape.
Here are some details about my environment:
Python: 3.8.20
Pytorch: 2.2.0
swarms-torch: 0.2.2
CUDA: 12.2
Could you please assist in resolving this issue or provide guidance on how to fix it?

Upvote & Fund

• We're using Polar.sh so you can upvote and help fund this issue.
• We receive the funding once the issue is completed & confirmed by you.
• Thank you in advance for helping prioritize & fund our backlog.