refactor llama/model.py

4 months ago · a945711f58
parent 3b00036ec6
commit a945711f58
1 changed files with 2 additions and 1 deletions
--- a/src/petals/models/llama/model.py
+++ b/src/petals/models/llama/model.py
@ -90,7 +90,8 @@ class DistributedLlamaModel(FromPretrainedMixin, PTuneMixin, LlamaModel):
            hypo_ids=past_key_values.hypo_ids if past_key_values is not None else None,
        )

-        past_key_values = past_key_values if past_key_values is not None else RemotePastKeyValues()
+        if past_key_values is None:
+            past_key_values = RemotePastKeyValues()
        past_key_values.update_seen(hidden_states.size(1))

        # Remove prefix