Add a repro for #3479

NVIDIA · Nov 26, 2024 · 22a05d5 · 22a05d5
1 parent 9c9c34c
commit 22a05d5
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/tests/cpp/test_allocation_domain.cpp b/tests/cpp/test_allocation_domain.cpp
@@ -1426,8 +1426,9 @@ TEST_F(AllocationDomainTest, InputAllocationIsSplit_Concrete) {
   fusion->addInput(in);
   fusion->addOutput(out);
 
-  in->split(0, 2);
-  in->setAllocationDomain(in->getLoopDomain(), true);
+  auto [outer, inner] = IterDomain::split(
+      in->axis(0), IrBuilder::create<Val>(2, DataType::Index), true);
+  in->setAllocationDomain({outer, inner}, true);
 
   FusionExecutorCache executor_cache(std::move(fusion));
   auto options = at::TensorOptions().dtype(at::kFloat).device(at::kCUDA);