File tree Expand file tree Collapse file tree 1 file changed +0
-13
lines changed
torchtitan/experiments/compiler_toolkit Expand file tree Collapse file tree 1 file changed +0
-13
lines changed Original file line number Diff line number Diff line change 2121from 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-
3724def export_joint (
3825 model , args , kwargs = None
3926) -> tuple [JointWithDescriptors , TracingContext ]:
You can’t perform that action at this time.
0 commit comments