fix test

tests/generation/test_utils.py

@@ -3139,8 +3139,8 @@ def test_model_kwarg_assisted_decoding_encoder_decoder(self):
         # PT-only test: TF doesn't support assisted decoding yet.
         # Bart subclass with a kwarg that distorts the output
         class FakeBart(BartForConditionalGeneration):
-            def forward(self, input_ids, foo=False, **kwargs):
-                outs = super().forward(input_ids, **kwargs)
+            def forward(self, input_ids, past_key_values, foo=False, **kwargs):
+                outs = super().forward(input_ids, past_key_values=past_key_values, **kwargs)
                 if foo:
                     outs["logits"][:, :, :] = 0.0
                 return outs
@@ -3170,9 +3170,9 @@ def prepare_inputs_for_generation(self, *args, foo=False, encoder_outputs=None,
             self.assertListEqual(outputs_foo.tolist(), outputs_normal.tolist())
 
         # Assistant model
-        assistant = AutoModelForSeq2SeqLM.from_pretrained(
-            "hf-internal-testing/tiny-random-BartForConditionalGeneration"
-        ).to(torch_device)
+        assistant = FakeBart.from_pretrained("hf-internal-testing/tiny-random-BartForConditionalGeneration").to(
+            torch_device
+        )
 
         # If assisted generation passes model_kwargs correctly, should be same as previous
         outputs_assisted = model.generate(