calibration_seq_length

pytorch · Nov 20, 2024 · 0019456 · 0019456
1 parent be822ae
commit 0019456
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/torchao/_models/llama/generate.py b/torchao/_models/llama/generate.py
@@ -312,6 +312,8 @@ def main(
             from torchao._models._eval import InputRecorder
             from torchao._models.llama.model import prepare_inputs_for_model
 
+            calibration_seq_length = 256
+            calibration_limit = 1
             inputs = InputRecorder(
                 tokenizer,
                 calibration_seq_length,
@@ -354,6 +356,8 @@ def main(
             from torchao._models._eval import InputRecorder
             from torchao._models.llama.model import prepare_inputs_for_model
 
+            calibration_seq_length = 256
+            calibration_limit = 1
             inputs = InputRecorder(
                 tokenizer,
                 calibration_seq_length,