adjust batch size for mistral-7b

plutonium-239 · May 21, 2024 · c0ad551 · c0ad551
1 parent 6f86fb5
commit c0ad551
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/experiments/paper_demo.py b/experiments/paper_demo.py
@@ -126,6 +126,8 @@
         for model in models:
             if model == "flan-t5":
                 batch_size = 56
+            if model == "mistral-7b":
+                batch_size = 16
             for estimate in estimators:
                 outputs = []