make func

Signed-off-by: Mayank Mishra <[email protected]>
IBM · Sep 12, 2024 · 94a1b46 · 94a1b46
1 parent b2c19d7
commit 94a1b46
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/dolomite_engine/hf_models/mixins/moe/base.py b/dolomite_engine/hf_models/mixins/moe/base.py
@@ -83,13 +83,14 @@ def forward(
 
             hidden_states = outputs[0]
             outputs = outputs[1:]
+
             if output_router_logits:
                 all_router_logits += (outputs[0],)
                 outputs = outputs[1:]
 
             if output_aux_loss:
-                total_aux_loss = total_aux_loss + outputs[0]
-                outputs = outputs[1:]
+                aux_loss = outputs[0]
+                total_aux_loss = total_aux_loss + aux_loss
 
         hidden_states = self.ln_f(hidden_states)