Fix

Author: mseeger

litgpt/generate/base.py

@@ -355,7 +355,7 @@ def batched_generate_fn(
                 input_pos=input_pos,
                 kwargs=sample_args,
             )
-        input_pos += 1
+            input_pos += 1
         for i in range(batch_size):
             token_lists[i].append(tokens[i])
         int_tokens = [token.item() for token in tokens]

tests/test_batch.py

@@ -244,7 +244,6 @@ def find_unique_stop(triplets):
         stop_tokens=tuple([s] for s in stops),
         sample_args=sample_kwargs,
         include_prompt=True,
-        include_eos=False,
     ):
         tokens.append([t.item() if t is not None else None for t in l])