Fix memory allocation of ndarray

abetlen · Aug 25, 2024 · f993209 · f993209
1 parent 259ee15
commit f993209
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/llama_cpp/llama.py b/llama_cpp/llama.py
@@ -2049,7 +2049,8 @@ def load_state(self, state: LlamaState) -> None:
         assert self._ctx.ctx is not None
         # Only filling in up to `n_tokens` and then zero-ing out the rest
         self.scores[: state.n_tokens, :] = state.scores.copy()
-        self.scores[state.n_tokens :, :] = 0.0
+        rest = self.scores[state.n_tokens :, :]
+        rest[rest > 0] = 0.0
         self.input_ids = state.input_ids.copy()
         self.n_tokens = state.n_tokens
         state_size = state.llama_state_size