From ae1bdc1a0b9cdc36137d4a719e2f0dffba5af32b Mon Sep 17 00:00:00 2001 From: Vasiliy Kuznetsov Date: Tue, 5 Jan 2021 10:52:49 -0800 Subject: [PATCH] eager quant: fix error with removing forward hooks (#49813) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/49813 https://github.com/pytorch/pytorch/issues/49739 reports a crash where removing forward hooks results in a ``` RuntimeError: OrderedDict mutated during iteration ``` Unfortunately I cannot repro this inside the PyTorch module, but the issue author has a good point and and we should not mutate the dict inside of the iteration. Test Plan: ``` // test plan from https://github.com/pytorch/pytorch/pull/46871 which // originally added this python test/test_quantization.py TestEagerModeQATOps ``` Imported from OSS Reviewed By: jerryzh168 Differential Revision: D25698725 fbshipit-source-id: 13069d0d5017a84038c8f7be439a3ed537938ac6 --- torch/quantization/quantize.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/torch/quantization/quantize.py b/torch/quantization/quantize.py index a57a4ea6bcb8..77752a8af9c9 100644 --- a/torch/quantization/quantize.py +++ b/torch/quantization/quantize.py @@ -258,9 +258,12 @@ def _remove_activation_post_process(module): delattr(module, 'activation_post_process') # remove activation_post_proceess hook + handle_ids_to_remove = set() for handle_id, hook_fn in module._forward_hooks.items(): if hook_fn is _observer_forward_hook: - module._forward_hooks.pop(handle_id) + handle_ids_to_remove.add(handle_id) + for handle_id in handle_ids_to_remove: + module._forward_hooks.pop(handle_id) # TODO: rename to something more general def _remove_qconfig(module):