adds str information

Lightning-AI · Feb 3, 2025 · cfe3ded · cfe3ded
1 parent a03f073
commit cfe3ded
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 1 deletion.
diff --git a/thunder/dynamo/report.py b/thunder/dynamo/report.py
@@ -396,6 +396,9 @@ def __init__(self, graphs: list[torch.fx.GraphModule], graph_names: list[str] =
             graph_names = [f"graph{idx}" for idx in range(len(graphs))]
         self.fx_graph_reports: list[FXGraphReport] = [FXGraphReport(g, name) for name, g in zip(graph_names, graphs)]
 
+    def __str__(self):
+        return f"<FXReport with {len(self.fx_graph_reports)} FXGraphReports accessible via .fx_graph_reports>"
+
 
 def fx_report(fn: Callable, *args, compile_options: dict = None, **kwargs) -> FXReport:
     """

diff --git a/thunder/tests/test_dynamo.py b/thunder/tests/test_dynamo.py
@@ -1117,7 +1117,7 @@ def check(file_name, cmd):
 
     cmd = [sys.executable]
     if use_benchmark:
-        cmd = cmd + ["-m", "pytest", "--benchmark-warmup-iterations=5"]
+        cmd = cmd + ["-m", "pytest"]
     py_files = list(tmp_path.glob("*.py"))
     assert len(py_files) == 4