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
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)
Pay now to fund the work behind this issue.
Get updates on progress being made.
Maintainer is rewarded once the issue is completed.
You're funding impactful open source efforts
You want to contribute to this effort
You want to get funding like this too