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