Fix test_circular_buffering.cpp (#3285)

NVIDIA · Oct 27, 2024 · 374fa38 · 374fa38
1 parent 5776198
commit 374fa38
Showing 1 changed file with 2 additions and 4 deletions.
diff --git a/tests/cpp/test_circular_buffering.cpp b/tests/cpp/test_circular_buffering.cpp
@@ -803,10 +803,9 @@ TEST_P(CircularBufferingTest, NoSync) {
   testValidate(&fusion, cg_outputs, {t0, t1}, {ref}, __LINE__, __FILE__);
 }
 
-// Test circular buffer from 2 to 10 stages
 auto StagesAndPrefetches() {
   std::vector<StageAndPrefetch> values;
-  for (int64_t i : {2, 5, 10}) {
+  for (int64_t i : {2, 5, 9}) {
     for (int64_t j : c10::irange(-i, i)) {
       values.emplace_back(i, j);
     }
@@ -1765,10 +1764,9 @@ TEST_P(TmaCircularBufferingTest, MatmulWithBroadcastedInput) {
       fusion.get(), cg_outputs, {t0, t1}, {aten_output}, __LINE__, __FILE__);
 }
 
-// Test circular buffer from 2 to 5 stages
 auto tmaCircularBufferingParams() {
   std::vector<TmaCircularBufferingParams> values;
-  for (int64_t i : {2, 5}) {
+  for (int64_t i : {2, 4}) {
     for (int64_t j : c10::irange(-i, i)) {
       for (int64_t m : {128, 500, 1024}) {
         for (int64_t n : {128, 1024}) {