main() — pytorch Function Reference

Source Code

benchmarks/dynamo/microbenchmarks/nested_graph_breaks.py lines 43–66

def main():
    start = time.perf_counter()
    print(top_fn(torch.ones(3)))
    end = time.perf_counter()

    print(f"eager total time: {end - start:.2f}s")

    opt_fn = torch.compile(top_fn, backend="eager")

    torch._dynamo.config.nested_graph_breaks = True
    start = time.perf_counter()
    opt_fn(torch.ones(3))
    end = time.perf_counter()

    print(f"nested_graph_breaks=True total time: {end - start:.2f}s")

    torch.compiler.reset()

    torch._dynamo.config.nested_graph_breaks = False
    start = time.perf_counter()
    opt_fn(torch.ones(3))
    end = time.perf_counter()

    print(f"nested_graph_breaks=False total time: {end - start:.2f}s")