Enable quantization of tied embeddings (#1703)

neuralmagic · Aug 14, 2023 · bb5021b · bb5021b
1 parent 45f051a
commit bb5021b
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/src/sparseml/pytorch/sparsification/quantization/helpers.py b/src/sparseml/pytorch/sparsification/quantization/helpers.py
@@ -706,7 +706,7 @@ def prepare_embeddings_qat(
     for submodule in module.modules():
         submodule_qconfig = getattr(submodule, "qconfig", None)
         submodule_qconfig = submodule_qconfig or qconfig
-        if type(submodule) is Embedding and submodule_qconfig is not None:
+        if isinstance(submodule, Embedding) and submodule_qconfig is not None:
             _prepare_qat_embedding(submodule, submodule_qconfig)