Fix CANINE fp16 dtype mismatch in attention mask causing ONNX export failure

src/transformers/models/canine/modeling_canine.py

@@ -341,7 +341,7 @@ def forward(
                 # positions we want to attend and the dtype's smallest value for masked positions.
                 attention_mask = (1.0 - attention_mask.float()) * torch.finfo(attention_scores.dtype).min
             # Apply the attention mask (precomputed for all layers in CanineModel forward() function)
-            attention_scores = attention_scores + attention_mask
+            attention_scores = attention_scores + attention_mask.to(attention_scores.dtype)
 
         # Normalize the attention scores to probabilities.
         attention_probs = nn.functional.softmax(attention_scores, dim=-1)

tests/models/canine/test_modeling_canine.py

@@ -248,6 +248,22 @@ def test_model(self):
         config_and_inputs = self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_model(*config_and_inputs)
 
+    def test_model_fp16_with_attention_mask(self):
+        config, input_ids, token_type_ids, input_mask, _, _, _ = self.model_tester.prepare_config_and_inputs()
+        # Force partial masking so the ndim==3 mask path in the shallow encoder is exercised.
+        input_mask[:, self.model_tester.seq_length // 2 :] = 0
+
+        model = CanineModel(config=config)
+        model.to(torch_device)
+        model.half()
+        model.eval()
+
+        with torch.no_grad():
+            output = model(input_ids, attention_mask=input_mask).last_hidden_state
+
+        self.assertEqual(output.dtype, torch.float16)
+        self.assertFalse(torch.isnan(output).any())
+
     def test_for_multiple_choice(self):
         config_and_inputs = self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_for_multiple_choice(*config_and_inputs)