add import torch

Signed-off-by: youkaichao <[email protected]>
thuml · Dec 6, 2024 · 810a8fb · 810a8fb
1 parent 87f5c38
commit 810a8fb
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/depyf/explain/patched_lazy_format_graph_code.py b/depyf/explain/patched_lazy_format_graph_code.py
@@ -27,7 +27,7 @@ def patched_lazy_format_graph_code(name, gm, maybe_id=None, **kwargs):
     use_gm = True
 
     # use `print_readable` because it can include submodules
-    src = "from __future__ import annotations\n" + \
+    src = "from __future__ import annotations\nimport torch\n" + \
         gm.print_readable(print_output=False)
     src = src.replace("<lambda>", "GraphModule")
     try: