Skip to content

Commit edbf349

Browse files
authored
[easy] [compiler toolkit] Clean up unused function (#2014)
As titled. `_clear_traced_params_buffers` is no longer being used as we have switched the dynamo graph capture API.
1 parent f5d2b18 commit edbf349

File tree

1 file changed

+0
-13
lines changed

1 file changed

+0
-13
lines changed

torchtitan/experiments/compiler_toolkit/graph_utils.py

Lines changed: 0 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -21,19 +21,6 @@
2121
from torchtitan.tools.logging import logger
2222

2323

24-
def _clear_traced_params_buffers(
25-
traced_module: torch.fx.GraphModule, const_keys: list[str]
26-
) -> None:
27-
"""Remove all parameters and buffers from traced module before restoring."""
28-
for key in const_keys:
29-
assert key in traced_module._buffers.keys()
30-
# We don't want constants to show up as a buffer in the state dict.
31-
# Instead they should just be a direct attribute.
32-
buffer = getattr(traced_module, key)
33-
torch.fx.graph_module._del_attr(traced_module, key)
34-
setattr(traced_module, key, buffer)
35-
36-
3724
def export_joint(
3825
model, args, kwargs=None
3926
) -> tuple[JointWithDescriptors, TracingContext]:

0 commit comments

Comments
 (0)