Merge branch 'merge-wilds-mllm' into mllm-eval-hl

mlfoundations · Nov 29, 2023 · 5d10063 · 5d10063
2 parents 55a8f7c + aa995e3
commit 5d10063
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()
@@ -421,7 +422,7 @@ def _prepare_inputs_for_forward(
                     v.repeat_interleave(num_beams, dim=0)
                 )
                 for k, v in past_key_values
-            ]
+
         return {
             "input_ids": lang_x,
             "attention_mask": attention_mask,