Fix test

huggingface · Dec 28, 2023 · 4320b1c · 4320b1c
1 parent 41164bb
commit 4320b1c
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/tests/test_tokenization_common.py b/tests/test_tokenization_common.py
@@ -2624,7 +2624,8 @@ def test_prepare_seq2seq_batch(self):
                     src_texts=src_text, max_length=3, max_target_length=10, return_tensors="pt"
                 )
                 self.assertEqual(batch_encoder_only.input_ids.shape[1], 3)
-                self.assertEqual(batch_encoder_only.attention_mask.shape[1], 3)
+                if "attention_mask" in tokenizer.model_input_names:
+                    self.assertEqual(batch_encoder_only.attention_mask.shape[1], 3)
                 self.assertNotIn("decoder_input_ids", batch_encoder_only)
 
     def test_is_fast(self):