From d280339a67e30536f0faa5310b79b71d5a5a1d68 Mon Sep 17 00:00:00 2001 From: Ubuntu Date: Wed, 10 Jun 2020 14:36:01 +0000 Subject: [PATCH] fix calibration pass to support multiple functions --- 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 59ee51b7ef59..9590e87534d1 100644 --- a/python/tvm/relay/quantize/_calibrate.py +++ b/python/tvm/relay/quantize/_calibrate.py @@ -138,10 +138,14 @@ def _make_const(val): 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