Skip to content

Commit

Permalink
Comments
Browse files Browse the repository at this point in the history
  • Loading branch information
daniil-lyakhov committed Oct 15, 2024
1 parent 6b61fa2 commit c96c91a
Showing 1 changed file with 5 additions and 3 deletions.
8 changes: 5 additions & 3 deletions nncf/experimental/torch/fx/constant_folding.py
Original file line number Diff line number Diff line change
Expand Up @@ -240,15 +240,17 @@ def insert_placerholder_values(self, env: Dict[torch.fx.Node, Any]) -> None:

def constant_fold(
gm: torch.fx.GraphModule,
constraint_fn: Optional[Callable[[torch.fx.Node], bool]] = None,
) -> None:
"""
Calcualtes constant subgraphs values and replaces them with a constan node inplace.
:param gm: Given graph model.
"""
with torch.utils._python_dispatch._disable_current_modes():
cf = ConstantFolder(gm, skip_constructors=True)
cf.run()

for node, constant in cf.node_replacements.items():
if constraint_fn is not None and not constraint_fn(node):
continue
replace_node_with_constant(gm, node, constant)

erased_params = []
Expand Down

0 comments on commit c96c91a

Please sign in to comment.