From af8f5a0f2ab4034501e4e6b37d5e90a9002ea040 Mon Sep 17 00:00:00 2001
From: mgoin <michael@neuralmagic.com>
Date: Thu, 18 Jul 2024 17:54:06 -0400
Subject: [PATCH] Fix test

---
 auto_fp8/modeling.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/auto_fp8/modeling.py b/auto_fp8/modeling.py
index eb4d2ba..0e4e8cc 100644
--- 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):