From 15d1bd4d1f77b3dfe58e453ecbe28bbe5ae8eab9 Mon Sep 17 00:00:00 2001 From: Megvii Engine Team Date: Mon, 1 Jun 2020 15:58:54 +0800 Subject: [PATCH] fix(mge/quantization): use list rather than generator to modify submodules GitOrigin-RevId: b5a3851330e7ffeb2aa9ad1a42e8ade8a0817506 --- python_module/megengine/quantization/quantize.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/python_module/megengine/quantization/quantize.py b/python_module/megengine/quantization/quantize.py index 2424e22ff..358c6cc50 100644 --- a/python_module/megengine/quantization/quantize.py +++ b/python_module/megengine/quantization/quantize.py @@ -65,9 +65,9 @@ def quantize(module: Module, inplace=True): def is_qat(mod: Module): return isinstance(mod, qat_modules) - # no need to pass prefix and get pure key of parent Module. - for key, submodule, parent in module._flatten( - with_key=True, with_parent=True, predicate=is_qat + # must use list to avoid replacement influencing successor modules + for key, submodule, parent in list( + module._flatten(with_key=True, with_parent=True, predicate=is_qat) ): new_mod = _qat2quantized_dict[type(submodule)].from_qat_module(submodule) if isinstance(parent, Float.Sequential): @@ -100,9 +100,9 @@ def quantize_qat( def is_quantable(mod: Module): return isinstance(mod, quantable_modules) - # no need to pass prefix and get pure key of parent Module. - for key, submodule, parent in module._flatten( - with_key=True, with_parent=True, predicate=is_quantable + # must use list to avoid replacement influencing successor modules + for key, submodule, parent in list( + module._flatten(with_key=True, with_parent=True, predicate=is_quantable) ): # only convert top quantable module. if is_quantable(parent): -- GitLab