Fix test

neuralmagic · Jul 18, 2024 · af8f5a0 · af8f5a0
1 parent b4f830d
commit af8f5a0
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/auto_fp8/modeling.py b/auto_fp8/modeling.py
@@ -79,6 +79,7 @@ def quantize(self, dataset: Optional[Dataset] = None):
             model=self.model,
             dataset=dataset,
             recipe=recipe,
+            num_calibration_samples=dataset.shape[0],
         )
 
     def save_quantized(self, save_directory: str):