From d7bcedf421d0aa643ab05549553965d7ff007577 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E5=91=A8=E5=91=A8=E5=91=A8?=
 <39978853+zhoutianzi666@users.noreply.github.com>
Date: Wed, 22 Oct 2025 17:49:18 +0800
Subject: [PATCH] small change in test_fusedmoe.py (#4538)

---
 tests/layers/test_fusedmoe.py | 13 +++++++++----
 1 file changed, 9 insertions(+), 4 deletions(-)

diff --git a/tests/layers/test_fusedmoe.py b/tests/layers/test_fusedmoe.py
index 59d7d30f2..521005090 100644
--- a/tests/layers/test_fusedmoe.py
+++ b/tests/layers/test_fusedmoe.py
@@ -185,12 +185,18 @@ class TestFusedMoE(unittest.TestCase):
 
             cache_hidden_states[idx] = paddle.rand((num_tokens, self.model_config.hidden_size), dtype=paddle.bfloat16)
 
+            num_layers = 80
+
+            def fake_model_run():
+                for _ in range(num_layers):
+                    out = fused_moe.fused_moe(cache_hidden_states[idx], gating)
+
+                return out
+
             moe_cuda_graphs[idx] = graphs.CUDAGraph()
             moe_cuda_graphs[idx].capture_begin()
 
-            num_layers = 80
-            for _ in range(num_layers):
-                out = fused_moe.fused_moe(cache_hidden_states[idx], gating)
+            fake_model_run()
 
             moe_cuda_graphs[idx].capture_end()
 
@@ -213,7 +219,6 @@ class TestFusedMoE(unittest.TestCase):
             print(times[-1], round(GB / times_s, 1))
 
         shutil.rmtree(self.model_name_or_path)
-        return out
 
 
 if __name__ == "__main__":