missing device (#232)

pytorch · Jul 17, 2024 · 4cda110 · 4cda110
1 parent 6dfce10
commit 4cda110
Showing 1 changed file with 6 additions and 2 deletions.
diff --git a/quantize.py b/quantize.py
@@ -515,7 +515,11 @@ def forward(self, input: torch.Tensor) -> torch.Tensor:
 
 
 def replace_embedding_weight_only_grouped_int8_per_channel(
-    module, device, bitwidth: int = 8, groupsize: Optional[int] = None, packed=False
+        module,
+        device,
+        bitwidth: int = 8,
+        groupsize: Optional[int] = None,
+        packed=False
 ):
     for name, child in module.named_children():
         # print(f"name: {name}")
@@ -535,7 +539,7 @@ def replace_embedding_weight_only_grouped_int8_per_channel(
             )
         else:
             replace_embedding_weight_only_grouped_int8_per_channel(
-                child, bitwidth, groupsize, packed
+                child, device, bitwidth, groupsize, packed
             )