From 5ba766373261b8dca2d936e93395201d07c38d2a Mon Sep 17 00:00:00 2001 From: Megvii Engine Team Date: Wed, 22 Jun 2022 15:31:25 +0800 Subject: [PATCH] fix(traced_module): fix trace_module function may raise error in finally scope GitOrigin-RevId: 2689b1f6e6cc91c8c1f3f18e67e2860851e78263 --- imperative/python/megengine/traced_module/traced_module.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/imperative/python/megengine/traced_module/traced_module.py b/imperative/python/megengine/traced_module/traced_module.py index 00b7852b..0c33fe39 100644 --- a/imperative/python/megengine/traced_module/traced_module.py +++ b/imperative/python/megengine/traced_module/traced_module.py @@ -34,7 +34,7 @@ from ..core._imperative_rt.core2 import ( set_module_tracing, unset_module_tracing, ) -from ..core._trace_option import set_symbolic_shape +from ..core._trace_option import set_symbolic_shape, use_symbolic_shape from ..core.ops.builtin import Copy from ..module import Module from ..module import external as MExternal @@ -2451,6 +2451,8 @@ def trace_module( """ assert active_module_tracer() is None assert isinstance(mod, Module) + use_sym_shape = use_symbolic_shape() + inputs = [] try: net_name = mod._name if mod._name else mod.__class__.__name__ use_sym_shape = set_symbolic_shape(True)