Fixed __call__ as well

outlines-dev · Jul 8, 2024 · 937f0f4 · 937f0f4
1 parent 0800552
commit 937f0f4
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/outlines/generate/api.py b/outlines/generate/api.py
@@ -231,8 +231,8 @@ def __call__(
 
         # We reshape the output to (batch_size, sample_size)
         output: List[List[FormattedOutput]] = list()
-        for i in range(batch_size):
-            output.append(formatted[i : i + num_samples])
+        for i in range(0, batch_size * num_samples, num_samples):
+            output.append(formatted[i : i + num_samples]
 
         # We remove leading dimensions for the output
         if batch_size == 1 and num_samples == 1: