Update modeling_tiny_mixtral.py

modeling_tiny_mixtral.py

@@ -456,7 +456,7 @@ class SparseMOE(nn.Module):
         alpha = 0.01  # auxiliary loss weight (you can make this configurable)
         load_balancing_loss = alpha * self.num_experts * torch.sum(f_i * P_i)
         
-        return final_output, load_balancing_loss
+        return final_output
 
 ##final_loss = task_loss + router_loss_weight * router_loss