diff --git a/tests/e2e/offline_inference/test_omni_sleep_mode.py b/tests/e2e/offline_inference/test_omni_sleep_mode.py
index 2ad0b53b010..fc4001ff596 100644
--- a/tests/e2e/offline_inference/test_omni_sleep_mode.py
+++ b/tests/e2e/offline_inference/test_omni_sleep_mode.py
@@ -123,7 +123,7 @@ async def test_multistage_sleep_h100(tp_size):
             pass
 
         acks = await engine.sleep(stage_ids=[0, 1], level=2)
-        assert len(acks) == 2 * tp_size
+        assert len(acks) == 2
 
         await engine.wake_up(stage_ids=[0, 1])
         async for _ in engine.generate("verify", sampling_params=[SamplingParams(), sp]):