concatenation correction

This commit is contained in:
Nathan Ranchin
2024-09-10 22:05:03 +02:00
committed by Awni Hannun
parent 67a9325603
commit 92ba4b2bf5

View File

@@ -255,7 +255,7 @@ def generate_step(
else:
y, logprobs = sample(logits)
tokens_ids = mx.concat(tokens_ids, y, dim=0)
tokens_ids = mx.concat([tokens_ids, y], axis=0)
if repetition_context_size:
if len(repetition_context) > repetition_context_size: