fix

allenai · Oct 11, 2024 · 2ffa147 · 2ffa147
1 parent baa0f89
commit 2ffa147
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/src/olmo_core/nn/moe/layers.py b/src/olmo_core/nn/moe/layers.py
@@ -44,7 +44,7 @@ def get_loss(self) -> torch.Tensor:
             clear_router_zloss,
         )
 
-        _, expert_scores, _ = zip(*get_load_balancing_loss())
+        _, expert_scores = zip(*get_load_balancing_loss())
         tokens = expert_scores[0].shape[0]
 
         loss = batched_load_balancing_loss(self.args)