shape consistency

neuralmagic · May 1, 2024 · bf7d0f6 · bf7d0f6
1 parent c894305
commit bf7d0f6
Showing 1 changed file with 13 additions and 1 deletion.
diff --git a/src/sparseml/modifiers/obcq/utils/sgpt_wrapper.py b/src/sparseml/modifiers/obcq/utils/sgpt_wrapper.py
@@ -179,10 +179,22 @@ def fasterprune(
                             fake_quantize,
                         )
 
+                        while scale.ndim < 2:
+                            scale = scale.unsqueeze(1)
+                            zero_point = zero_point.unsqueeze(1)
+
+                        while q.ndim < 2:
+                            q = q.unsqueeze(1)
                         q = fake_quantize(
-                            q, scale, zero_point, self.layer.quantization_scheme.weights
+                            q,
+                            scale[:, i],
+                            zero_point[:, i],
+                            self.layer.quantization_scheme.weights,
                         )
 
+                while q.ndim != 1:
+                    q.squeeze()
+
                 Q1[:, i] = q
                 Losses1[:, i] = (w - q) ** 2 / d**2