Pass all arguments when quantizing embedding bag from float (#137697)

annop-w · pytorchmergebot · commit 81ecf98d23a7 · 2024-11-07T09:53:49.000Z
Pull Request resolved: #137697 Approved by: https://github.com/snadampal, https://github.com/jerryzh168
diff --git a/torch/ao/nn/quantized/modules/embedding_ops.py b/torch/ao/nn/quantized/modules/embedding_ops.py
@@ -383,7 +383,15 @@ def from_float(cls, mod, use_precomputed_fake_quant=False):
 
         # Create quantized EmbeddingBag module and pass in the quantized weight
         qembedding_bag = EmbeddingBag(
-            mod.num_embeddings, mod.embedding_dim, dtype=dtype
+            mod.num_embeddings,
+            mod.embedding_dim,
+            max_norm=mod.max_norm,
+            norm_type=mod.norm_type,
+            scale_grad_by_freq=mod.scale_grad_by_freq,
+            mode=mod.mode,
+            sparse=mod.sparse,
+            include_last_offset=mod.include_last_offset,
+            dtype=dtype,
         )
         qembedding_bag.set_weight(qweight)
         return qembedding_bag
