a bit more strict check

Signed-off-by: Masaki Kozuki <mkozuki@nvidia.com>
Lightning-AI · Jul 19, 2024 · a893e96 · a893e96
1 parent 54d886b
commit a893e96
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/thunder/tests/distributed/test_ops.py b/thunder/tests/distributed/test_ops.py
@@ -324,7 +324,7 @@ def foo(
             if bsym.sym.id in {all_gather_prim_impl.id, reduce_scatter_prim_impl.id}:
                 comm_idx = idx
             if bsym.sym.id == wait_prim_impl.id:
-                self.assertGreater(idx, comm_idx)
+                self.assertGreater(idx, comm_idx + 2)
 
 
 common_utils.instantiate_parametrized_tests(DistributedCollectiveOpTest)