[CI] fix patch_observer_set for PT 1.9 (#321)

alibaba · May 31, 2024 · 1e74c0a · 1e74c0a
1 parent 8c1f2ce
commit 1e74c0a
Showing 1 changed file with 6 additions and 1 deletion.
diff --git a/tinynn/graph/quantization/quantizer.py b/tinynn/graph/quantization/quantizer.py
@@ -3711,11 +3711,16 @@ def new_no_observer_set():
 
                     return new_no_observer_set
 
-                if LooseVersion(torch.__version__) >= LooseVersion("1.9.0"):
+                if LooseVersion(torch.__version__) >= LooseVersion("1.10.0"):
                     orig_no_observer_set = sys.modules['torch.ao.quantization.quantize'].no_observer_set
                     sys.modules['torch.ao.quantization.quantize'].no_observer_set = patch_observer_set(
                         orig_no_observer_set
                     )
+                elif LooseVersion(torch.__version__) >= LooseVersion("1.9.0"):
+                    orig_no_observer_set = torch.ao.quantization.quantization_mappings.no_observer_set
+                    torch.ao.quantization.quantization_mappings.no_observer_set = patch_observer_set(
+                        orig_no_observer_set
+                    )
 
                 add_observer_func(
                     graph.module,