fixing import bug and gating by torch version

Summary: caused #87 when trying to migrate fvcore to ao, this fixes by gating before importing Differential Revision: D31706627 fbshipit-source-id: 1b0374325093e8492bb93afd7ea8e8b140d8bcb6
facebookresearch · Oct 16, 2021 · 21171ec · 21171ec
1 parent 4a39fce
commit 21171ec
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 2 deletions.
diff --git a/fvcore/common/checkpoint.py b/fvcore/common/checkpoint.py
@@ -18,7 +18,11 @@
 if TORCH_VERSION >= (1, 11):
     from torch.ao import quantization
     from torch.ao.quantization import ObserverBase, FakeQuantizeBase
-else:
+elif (
+    TORCH_VERSION >= (1, 8)
+    and hasattr(torch.quantization, "FakeQuantizeBase")
+    and hasattr(torch.quantization, "ObserverBase")
+):
     from torch import quantization
     from torch.quantization import ObserverBase, FakeQuantizeBase
 

diff --git a/tests/test_checkpoint.py b/tests/test_checkpoint.py
@@ -24,7 +24,11 @@
         disable_observer,
         enable_fake_quant,
     )
-else:
+elif (
+    TORCH_VERSION >= (1, 8)
+    and hasattr(torch.quantization, "FakeQuantizeBase")
+    and hasattr(torch.quantization, "ObserverBase")
+):
     from torch import quantization
     from torch.quantization import (
         get_default_qat_qconfig,