fix(partitioner): Update adjacency partitioner to handle changes in

torch.fx.passes.splitter_base._SplitterBase Signed-off-by: Naren Dasan <[email protected]> Signed-off-by: Naren Dasan <[email protected]>
pytorch · Apr 16, 2024 · b21fb5f · b21fb5f
1 parent 6dd6915
commit b21fb5f
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/py/torch_tensorrt/dynamo/partitioning/_adjacency_partitioner.py b/py/torch_tensorrt/dynamo/partitioning/_adjacency_partitioner.py
@@ -110,6 +110,7 @@ def __init__(
         allowed_single_node_partition_ops: Optional[Collection[str]] = None,
         min_block_size: int = MIN_BLOCK_SIZE,
         require_full_compilation: bool = REQUIRE_FULL_COMPILATION,
+        return_tuple: bool = False,
     ):
         """
         Preprocesses graph before splitting:
@@ -149,6 +150,7 @@ def __init__(
         self.num_trt_accelerated_subgraphs: Optional[int] = None
         self.allowed_single_node_partition_ops = allowed_single_node_partition_ops
         self.require_full_compilation = require_full_compilation
+        self._return_tuple = return_tuple
 
     def remove_small_acc_subgraphs(self, subgraphs: List[Subgraph]) -> List[Subgraph]:
         """