OSSForks/pytorch
1
0
Fork 0
mirror of https://github.com/zebrajr/pytorch.git synced 2025-12-06 12:20:52 +01:00
Code Issues Actions 38 Packages Projects Releases Wiki Activity

[Sigmoid] Remove magic method in CapabilityBasedPartitioner (#149400)

Browse Source 
Summary: As title.

Test Plan: CI

Differential Revision: D70575197

Pull Request resolved: https://github.com/pytorch/pytorch/pull/149400
Approved by: https://github.com/jfix71
This commit is contained in:
Ze Sheng 2025-03-19 16:02:40 +00:00 committed by PyTorch MergeBot
parent 4df66e0b7f
commit e98afa0f89
1 changed files with 2 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
torch/fx/passes/infra/partitioner.py
View File

@ -69,7 +69,7 @@ class CapabilityBasedPartitioner:
)
self.dependency_viewer = _DependencyViewer(graph_module)
def __is_node_supported(self, node: Node) -> bool:
def _is_node_supported(self, node: Node) -> bool:
return self.operator_support.is_node_supported(
dict(self.graph_module.named_modules()), node
)
@ -206,7 +206,7 @@ class CapabilityBasedPartitioner:
#
# I don't see a need to add a knob to disable horizontal fusion yet, we can short-cut
# the fusion by adding an `else` block here to skip horizontal fusion.
if self.__is_node_supported(node) and node not in assignment:
if self._is_node_supported(node) and node not in assignment:
partition_id = next(new_partition_id)
nodes_order[node] = partition_id
partitions_order[partition_id] = partition_id