Generate: add missing **model_kwargs in sample tests (#18696)

2022-08-19 16:14:27 +01:00
parent e54a1b49aa
commit e95d433d77
1 changed files with 2 additions and 0 deletions
--- a/tests/generation/test_generation_utils.py
+++ b/tests/generation/test_generation_utils.py
@@ -327,6 +327,7 @@ class GenerationTesterMixin:
            remove_invalid_values=True,
            **logits_warper_kwargs,
            **process_kwargs,
+            **model_kwargs,
        )

        torch.manual_seed(0)
@@ -361,6 +362,7 @@ class GenerationTesterMixin:
                **kwargs,
                **model_kwargs,
            )
+
        return output_sample, output_generate

    def _beam_search_generate(