diff --git a/userbenchmark/dynamo/dynamobench/_dynamo/utils.py b/userbenchmark/dynamo/dynamobench/_dynamo/utils.py index 809adc6b04..a7ce87277f 100644 --- a/userbenchmark/dynamo/dynamobench/_dynamo/utils.py +++ b/userbenchmark/dynamo/dynamobench/_dynamo/utils.py @@ -97,6 +97,7 @@ from torch._dispatch.python import enable_python_dispatcher from torch._utils_internal import log_compilation_event +from torch.fx._utils import _format_graph_code, lazy_format_graph_code from torch.nn.modules.lazy import LazyModuleMixin from torch.utils._triton import has_triton, has_triton_package @@ -1996,26 +1997,6 @@ def tensor_always_has_static_shape( return False, None -def lazy_format_graph_code(name, gm, maybe_id=None): - def format_name(): - if maybe_id is not None: - return f"{name} {maybe_id}" - else: - return name - - return LazyString( - lambda: _format_graph_code( - f"===== {format_name()} =====\n", - gm.forward.__code__.co_filename, - gm.print_readable(print_output=False), - ) - ) - - -def _format_graph_code(name, filename, graph_str): - return f"TRACED GRAPH\n {name} {filename} {graph_str}\n" - - def lazy_format_graph_tabular(fn_name, gm): def inner(): try: