added clone to fix prefill issue

pytorch-labs · May 5, 2024 · 2c33914 · 2c33914
1 parent 30d69b3
commit 2c33914
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/generate.py b/generate.py
@@ -172,7 +172,7 @@ def generate(
     seq = empty
     input_pos = torch.arange(0, T, device=device)
 
-    next_token = prefill(model, prompt.view(1, -1), input_pos, **sampling_kwargs)
+    next_token = prefill(model, prompt.view(1, -1), input_pos, **sampling_kwargs).clone()
     if is_speculative:
         prefill(draft_model, prompt.view(1, -1), input_pos, **sampling_kwargs)
     seq[T] = next_token