Update tensorrt_llm/_torch/modules/fused_moe/configurable_moe.py

Co-authored-by: Enwei Zhu <21126786+syuoni@users.noreply.github.com>
Signed-off-by: xxi <95731198+xxi-nv@users.noreply.github.com>

Author: xxi-nv

tensorrt_llm/_torch/modules/fused_moe/configurable_moe.py

@@ -440,7 +440,7 @@ def forward_impl(
             outputs = outputs[: all_rank_num_tokens[self.mapping.tp_rank]]
 
         # EPLB repeat logic
-        self.repeat_idx = 0 if self.repeat_idx == self.repeat_count - 1 else self.repeat_idx + 1
+        self.repeat_idx = (self.repeat_idx + 1) % self.repeat_count
 
         return outputs