OSSForks/pytorch
1
0
Fork 0
mirror of https://github.com/zebrajr/pytorch.git synced 2025-12-06 12:20:52 +01:00
Code Issues Actions 38 Packages Projects Releases Wiki Activity

[Easy] Use prepend arg to register forward hooks in quantize.py (#89391)

Browse Source 
Differential Revision: [D41431110](https://our.internmc.facebook.com/intern/diff/D41431110)
Pull Request resolved: https://github.com/pytorch/pytorch/pull/89391
Approved by: https://github.com/awgu
This commit is contained in:
Shen Li 2022-11-20 22:54:45 +00:00 committed by PyTorch MergeBot
parent 1db5ce095f
commit e0251de42f
1 changed files with 6 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

10
torch/ao/quantization/quantize.py
View File

@ -143,11 +143,13 @@ def register_activation_post_process_hook(module, pre_hook=False):
assert hasattr(module, 'activation_post_process'), \ assert hasattr(module, 'activation_post_process'), \
'Expect activation_post_process attribute already attached to the module' 'Expect activation_post_process attribute already attached to the module'
if pre_hook: if pre_hook:
handle = module.register_forward_pre_hook(_observer_forward_pre_hook) handle = module.register_forward_pre_hook(
module._forward_pre_hooks.move_to_end(handle.id, last=False) _observer_forward_pre_hook, prepend=True
)
else: else:
handle = module.register_forward_hook(_observer_forward_hook) handle = module.register_forward_hook(
module._forward_hooks.move_to_end(handle.id, last=False) _observer_forward_hook, prepend=True
)
def add_observer_(module, qconfig_propagation_list=None, non_leaf_module_list=None, device=None, custom_module_class_mapping=None): def add_observer_(module, qconfig_propagation_list=None, non_leaf_module_list=None, device=None, custom_module_class_mapping=None):