Some fixes in the recently added quantizer (#111)

Author: jerryzh168

test/quantization/model.py

@@ -12,6 +12,10 @@
 from torch.nn import functional as F
 
 def prepare_inputs_for_model(inps):
+    # this is because input from lm-eval is 2d
+    if input.dim() != 2:
+        raise ValueError(f"Expected input to be of dim 2, but got {input.dim()}")
+
     inps = inps.squeeze(0)
     # setup inputs in correct format
     max_new_tokens = 1

torchao/quantization/GPTQ.py

@@ -1225,9 +1225,9 @@ def __init__(
                 calibration_limit,
                 calibration_seq_length,
                 pad_calibration_inputs,
-                inner_k_tiles=8,
-                padding_allowed=True,
-                precision=torch.float32,
-                _is_gpt_fast=True,
+                inner_k_tiles=inner_k_tiles,
+                padding_allowed=padding_allowed,
+                precision=precision,
+                _is_gpt_fast=_is_gpt_fast,
                 _use_cuda=_use_cuda,
             )