Add generate kwargs to VQA pipeline (#29134)

a3f9221a · regisss · GitHub · 871ba71d · a3f9221a
Unverified Commit a3f9221a authored Feb 27, 2024 by regisss Committed by GitHub Feb 27, 2024
Show whitespace changes
Inline Side-by-side

Showing with 2 additions and 2 deletions

src/transformers/pipelines/visual_question_answering.py src/transformers/pipelines/visual_question_answering.py +2 -2

No files found.
--- a/src/transformers/pipelines/visual_question_answering.py
+++ b/src/transformers/pipelines/visual_question_answering.py
@@ -123,9 +123,9 @@ class VisualQuestionAnsweringPipeline(Pipeline):
        model_inputs.update(image_features)
        return model_inputs

-    def _forward(self, model_inputs):
+    def _forward(self, model_inputs, **generate_kwargs):
        if self.model.can_generate():
-            model_outputs = self.model.generate(**model_inputs)
+            model_outputs = self.model.generate(**model_inputs, **generate_kwargs)
        else:
            model_outputs = self.model(**model_inputs)
        return model_outputs