fix generation use_cache issue for mpt

mlfoundations · Nov 1, 2023 · aa995e3 · aa995e3
1 parent 17cd546
commit aa995e3
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/open_flamingo/src/vlm.py b/open_flamingo/src/vlm.py
@@ -277,6 +277,7 @@ def generate(
             **new_inputs,
             past_key_values=past_key_values,
             num_beams=num_beams,
+            use_cache=True,
             **kwargs,
         )
         self._post_forward_hook()
@@ -420,7 +421,7 @@ def _prepare_inputs_for_forward(
                 v.repeat_interleave(num_beams, dim=0)
             )
             for k, v in past_key_values
-        ]
+        ] if past_key_values is not None else None
         return {
             "input_ids": lang_x,
             "attention_mask": attention_mask,