From c70b27b9f7caede6de99373eb7e5aab7c2907df4 Mon Sep 17 00:00:00 2001 From: Denis Vieriu Date: Wed, 24 Apr 2024 13:01:27 -0700 Subject: [PATCH] Remove the sorting of the nodes from partitioning (not needed for now as Custom Metal kernels are yet not enabled) --- backends/apple/mps/partition/mps_partitioner.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/backends/apple/mps/partition/mps_partitioner.py b/backends/apple/mps/partition/mps_partitioner.py index 3dfc73cdd98..e5497389d14 100644 --- a/backends/apple/mps/partition/mps_partitioner.py +++ b/backends/apple/mps/partition/mps_partitioner.py @@ -102,7 +102,7 @@ def use_metal_kernel(self, node: torch.fx.Node): def tag_nodes(self, partitions: List[Partition]) -> None: for partition in partitions: crt_partition_counter = 0 - for node in sorted(partition.nodes): + for node in partition.nodes: delegation_tag = f"mps_{partition.id}" if self.use_metal_kernel(node): logging.warning(f"[WARNING] Using Metal kernel for op {node.name}!")