From 687d76579f2562135cf9aefa24d22be5d9a3835c Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E5=BC=A0=E7=A7=8B?= <zhangqiu@sensetime.com>
Date: Wed, 24 Jul 2024 17:10:13 +0800
Subject: [PATCH] fix eval for moe layer

---
 megablocks/layers/moe.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/megablocks/layers/moe.py b/megablocks/layers/moe.py
index 9d26da2e..264847ab 100644
--- a/megablocks/layers/moe.py
+++ b/megablocks/layers/moe.py
@@ -430,7 +430,7 @@ def forward(self, x, scores, expert_weights, top_experts):
         # Compute the experts.
         x, tokens_per_expert = self.forward_fn(
             x, expert_weights, top_experts)
-        if self.training and self.args.moe_loss_weight > 0:
+        if self.args.moe_loss_weight > 0:
             save_load_balancing_loss((tokens_per_expert, scores))
         x = x.view(in_shape)
         if self.bias is not None: