Update modeling_lfm2_moe_custom.py

modeling_lfm2_moe_custom.py

@@ -92,7 +92,7 @@ class Lfm2MoeDynamicTopKRouter(Lfm2MoeTopKRouter):
 
         self.budget = config.budget
         self.min_dynamic_k = 4
-        self.max_dynamic_k = min(32, self.num_experts)
+        self.max_dynamic_k = min(4, self.num_experts)
         self.layer_idx = layer_idx
 
         # ---- EXPERT ACTIVATION STATISTICS ----