[BUG] Dimension error within SwitchGate

[Masaaki-75]

Describe the bug
Shape mismatch is found in the computation of auxiliary loss values:

SwitchTransformers/switch_transformers/model.py

Lines 70 to 74 in 36a1ea0

	load = gate_scores.sum(0) # Sum over all examples
	importance = gate_scores.sum(1) # Sum over all experts
	# Aux loss is mean suqared difference between load and importance
	loss = ((load - importance) ** 2).mean()

where load is of shape [num_experts, dim] and importance is of shape [batch_size, dim]. Testing this SwitchGate class alone by giving an input with batch_size > 1 will raise error like this RuntimeError: The size of tensor a (64) must match the size of tensor b (2) at non-singleton dimension 0

To Reproduce
Simply run a sample with batch_size > 1:

gate = SwitchGate(dim=16, num_experts=3)
x = torch.randn((2, 64, 16)).float()
y, loss = gate(x, use_aux_loss=True)

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.

[github-actions]

Hello there, thank you for opening an Issue ! 🙏🏻 The team was notified and they will get back to you asap.

[github-actions]

Stale issue message