From 7e4428749c1dde026fb2cb4ea999f94d0755df30 Mon Sep 17 00:00:00 2001 From: Lysandre Debut Date: Thu, 11 Mar 2021 12:58:15 -0500 Subject: [PATCH] Conversion to tensors requires padding (#10661) --- tests/test_modeling_marian.py | 4 +++- tests/test_modeling_tf_marian.py | 2 +- 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/tests/test_modeling_marian.py b/tests/test_modeling_marian.py index 8e2b5fc513..191a48af8b 100644 --- a/tests/test_modeling_marian.py +++ b/tests/test_modeling_marian.py @@ -354,7 +354,9 @@ class MarianIntegrationTest(unittest.TestCase): self.assertListEqual(self.expected_text, generated_words) def translate_src_text(self, **tokenizer_kwargs): - model_inputs = self.tokenizer(self.src_text, return_tensors="pt", **tokenizer_kwargs).to(torch_device) + model_inputs = self.tokenizer(self.src_text, padding=True, return_tensors="pt", **tokenizer_kwargs).to( + torch_device + ) self.assertEqual(self.model.device, model_inputs.input_ids.device) generated_ids = self.model.generate( model_inputs.input_ids, attention_mask=model_inputs.attention_mask, num_beams=2, max_length=128 diff --git a/tests/test_modeling_tf_marian.py b/tests/test_modeling_tf_marian.py index 8000e41b5f..e4ccb28f00 100644 --- a/tests/test_modeling_tf_marian.py +++ b/tests/test_modeling_tf_marian.py @@ -363,7 +363,7 @@ class AbstractMarianIntegrationTest(unittest.TestCase): self.assertListEqual(self.expected_text, generated_words) def translate_src_text(self, **tokenizer_kwargs): - model_inputs = self.tokenizer(self.src_text, **tokenizer_kwargs, return_tensors="tf") + model_inputs = self.tokenizer(self.src_text, **tokenizer_kwargs, padding=True, return_tensors="tf") generated_ids = self.model.generate( model_inputs.input_ids, attention_mask=model_inputs.attention_mask, num_beams=2, max_length=128 )