Skip to content

Commit 563f2ff

Browse files
authored
fix failure of tests/models/shieldgemma2/test_modeling_shieldgemma2.p… (#42022)
* fix failure of tests/models/shieldgemma2/test_modeling_shieldgemma2.py::ShieldGemma2IntegrationTest::test_model Signed-off-by: Wang, Yi <yi.a.wang@intel.com> * add common kwargs in processing_shieldgemma2.py --------- Signed-off-by: Wang, Yi <yi.a.wang@intel.com>
1 parent 6f479d5 commit 563f2ff

File tree

1 file changed

+4
-0
lines changed

1 file changed

+4
-0
lines changed

src/transformers/models/shieldgemma2/processing_shieldgemma2.py

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -126,6 +126,10 @@ def __call__(
126126
if not self.chat_template:
127127
raise ValueError("ShieldGemma 2 requires the use of a specific chat template")
128128

129+
common_kwargs = kwargs.setdefault("common_kwargs", {})
130+
if "return_tensors" in kwargs:
131+
common_kwargs["return_tensors"] = kwargs.pop("return_tensors")
132+
129133
# Disable pan and scan
130134
images_kwargs = kwargs.setdefault("images_kwargs", {})
131135
if images_kwargs.get("do_pan_and_scan") is True:

0 commit comments

Comments
 (0)