From 53370b910302f951129d2fd8a9bfb9c0d3ef7a0a Mon Sep 17 00:00:00 2001 From: "Yi-Hsiang (Sean) Lai" Date: Fri, 12 Jun 2020 11:33:20 -0400 Subject: [PATCH] fix calibration pass to support multiple functions (#5768) Co-authored-by: Ubuntu --- python/tvm/relay/quantize/_calibrate.py | 12 ++++++++---- 1 file changed, 8 insertions(+), 4 deletions(-) diff --git a/python/tvm/relay/quantize/_calibrate.py b/python/tvm/relay/quantize/_calibrate.py index 59ee51b..9590e87 100644 --- a/python/tvm/relay/quantize/_calibrate.py +++ b/python/tvm/relay/quantize/_calibrate.py @@ -138,10 +138,14 @@ def _set_params(mod, input_scale_func, weight_scale_func): const_params[nclip_min] = _make_const(- (valid_range - 1)) const_params[nclip_max] = _make_const((valid_range - 1)) - func = mod['main'] - _analysis.post_order_visit(func, visit_func) - func = _expr.bind(func, const_params) - return IRModule.from_expr(func) + main_func = mod['main'] + _analysis.post_order_visit(main_func, visit_func) + main_func = _expr.bind(main_func, const_params) + func_dict = {} + for global_var, func in mod.functions.items(): + if global_var.name_hint != 'main': + func_dict[global_var] = func + return IRModule.from_expr(main_func, func_dict) # weight scale functions -- 2.7.4