fix export bug when multiple blocks in graph (#1162) (#1180)

* fix expport bug when multiple blocks in graph * fix * refactor code
PaddlePaddle · Dec 7, 2022 · 2f93fc4 · 2f93fc4
1 parent 72fbe8b
commit 2f93fc4
Show file tree

Hide file tree

Showing 3 changed files with 210 additions and 175 deletions.
diff --git a/visualdl/component/graph/exporter.py b/visualdl/component/graph/exporter.py
@@ -26,7 +26,8 @@ def translate_graph(model, input_spec, verbose=True):
     with tempfile.TemporaryDirectory() as tmp:
         model._full_name = '{}[{}]'.format(model.__class__.__name__, "model")
         create_opname_scope(model)
-        paddle.jit.save(model, os.path.join(tmp, 'temp'), input_spec)
+        model = paddle.jit.to_static(model, input_spec)
+        paddle.jit.save(model, os.path.join(tmp, 'temp'))
         model_data = open(os.path.join(tmp, 'temp.pdmodel'), 'rb').read()
         result = analyse_model(model_data)
     if verbose: