diff --git a/tests/test_modeling_bart.py b/tests/test_modeling_bart.py
index ada32ab647..2bcc2c4b45 100644
--- a/tests/test_modeling_bart.py
+++ b/tests/test_modeling_bart.py
@@ -314,7 +314,8 @@ class BartHeadTests(unittest.TestCase):
     @unittest.skipIf(torch_device == "cpu", "Cant do half precision")
     def test_generate_fp16(self):
         config, input_ids, batch_size = self._get_config_and_data(output_past=True)
-        attention_mask = input_ids.ne(1)
+        input_ids = input_ids.half()
+        attention_mask = input_ids.ne(1).to(torch_device).half()
         lm_model = BartForConditionalGeneration(config).eval().to(torch_device).half()
         lm_model.generate(input_ids, attention_mask=attention_mask)