diff --git a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredCommunicationPlanning.cpp b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredCommunicationPlanning.cpp
index d85b01d..5c77214 100644
--- a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredCommunicationPlanning.cpp
+++ b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredCommunicationPlanning.cpp
@@ -262,13 +262,12 @@ static void collectClosure(Value value, Block &body, const DeferredInputPlan &pl
 
 } // namespace
 
-bool isDeferredFragmentAssemblyInput(
-    Value input, const ComputeInstance &consumerInstance) {
+bool isDeferredFragmentAssemblyInput(Value input) {
   auto blueprint = input.getDefiningOp<SpatBlueprintOp>();
   if (!blueprint || blueprint.getMode() != "fragment_assembly")
     return false;
   return llvm::all_of(getBlueprintFragments(blueprint), [&](Value fragment) {
-    return getProducerValueRef(fragment, &consumerInstance).has_value();
+    return getProducerValueRef(fragment, nullptr).has_value();
   });
 }
 
@@ -279,7 +278,7 @@ LogicalResult prepareSingleCpuInput(OpBuilder &, Location loc, Value input, Bloc
                                     Value graphLane, Value scheduledGraphLane,
                                     DeferredInputPlan &plan) {
   plan = {graphInput, {}, {}, {}, graphLane, scheduledGraphLane, {}, {}, {}, {}, 1, nullptr};
-  if (isDeferredFragmentAssemblyInput(input, consumerInstance)) {
+  if (isDeferredFragmentAssemblyInput(input)) {
     plan.blueprint = input.getDefiningOp<SpatBlueprintOp>();
     plan.originalSources = getBlueprintFragments(plan.blueprint);
     return success();
@@ -310,7 +309,7 @@ LogicalResult prepareMultiCpuTupleInput(OpBuilder &, Location loc, Value input,
                                         DeferredInputPlan &plan) {
   const ComputeInstance &representative = tuple.instances.front();
   plan = {graphInput, {}, {}, {}, graphLane, scheduledGraphLane, scheduledLane, {}, {}, {}, 1, nullptr};
-  if (isDeferredFragmentAssemblyInput(input, representative)) {
+  if (isDeferredFragmentAssemblyInput(input)) {
     plan.blueprint = input.getDefiningOp<SpatBlueprintOp>();
     plan.originalSources = getBlueprintFragments(plan.blueprint);
     return success();
diff --git a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredCommunicationPlanning.hpp b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredCommunicationPlanning.hpp
index 8793847..ebf68ee 100644
--- a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredCommunicationPlanning.hpp
+++ b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredCommunicationPlanning.hpp
@@ -22,8 +22,7 @@ struct DeferredInputPlan {
   Block *scalarizedHoistBlock = nullptr;
 };
 
-bool isDeferredFragmentAssemblyInput(Value input,
-                                     const ComputeInstance &consumerInstance);
+bool isDeferredFragmentAssemblyInput(Value input);
 
 LogicalResult prepareSingleCpuInput(OpBuilder &builder, Location loc, Value input,
                                     BlockArgument graphInput,
diff --git a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/ScheduledComputePlanning.cpp b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/ScheduledComputePlanning.cpp
index b548d79..aab4beb 100644
--- a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/ScheduledComputePlanning.cpp
+++ b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/ScheduledComputePlanning.cpp
@@ -179,7 +179,7 @@ LogicalResult collectPeftClassOperandsAndResults(
       for (Value weight : getComputeInstanceWeights(instance))
         appendUnique(peftClassPlan.weights, weight);
       for (Value input : getComputeInstanceInputs(instance))
-        if (!getProducerValueRef(input, &instance) && !isDeferredFragmentAssemblyInput(input, instance))
+        if (!getProducerValueRef(input, &instance) && !isDeferredFragmentAssemblyInput(input))
           appendUnique(peftClassPlan.inputs, input);
     }
     return success();
@@ -222,7 +222,7 @@ LogicalResult collectPeftClassOperandsAndResults(
       for (Value weight : getComputeInstanceWeights(instance))
         appendUnique(peftClassPlan.weights, weight);
       for (Value input : getComputeInstanceInputs(instance))
-        if (!getProducerValueRef(input, &instance) && !isDeferredFragmentAssemblyInput(input, instance))
+        if (!getProducerValueRef(input, &instance) && !isDeferredFragmentAssemblyInput(input))
           appendUnique(peftClassPlan.inputs, input);
     }
   }