[quant][feature] Adding an ability to skip immediate observation

ghstack-source-id: 6b6ebbbe0dc165a8e7d151adaf328923129a9b1d Pull Request resolved: #49670
pytorch · Dec 21, 2020 · 3b51076 · 3b51076
1 parent 5c25f8f
commit 3b51076
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/torch/quantization/quantize.py b/torch/quantization/quantize.py
@@ -152,7 +152,8 @@ def insert_activation_post_process(m, special_act_post_process=None):
         elif needs_observation(child) and type(child) in custom_module_class_mapping:
             observed_child = custom_module_class_mapping[type(child)].from_float(child)
             setattr(module, name, observed_child)
-            insert_activation_post_process(observed_child)
+            if not getattr(custom_module_class_mapping[type(child)], 'do_not_observe', False):
+                insert_activation_post_process(observed_child)
         else:
             add_observer_(child, qconfig_propagation_list, non_leaf_module_list, device, custom_module_class_mapping)