[UDOP] Fix tests (#29573)

* Fix tests

* Fix tests

* Remove no_split_modules

tests/models/udop/test_modeling_udop.py

@@ -226,6 +226,20 @@ class UdopModelTester:
         )
         self.parent.assertTrue(torch.all(output_with_past_cache == output_without_past_cache))
 
+    def create_and_check_model_fp16_forward(
+        self,
+        config,
+        input_ids,
+        bbox,
+        decoder_input_ids,
+        attention_mask,
+        decoder_attention_mask,
+        lm_labels,
+    ):
+        model = UdopForConditionalGeneration(config=config).to(torch_device).half().eval()
+        output = model(input_ids, bbox=bbox, attention_mask=attention_mask, decoder_input_ids=decoder_input_ids).logits
+        self.parent.assertFalse(torch.isnan(output).any().item())
+
     def prepare_config_and_inputs_for_common(self):
         config_and_inputs = self.prepare_config_and_inputs()
         (
@@ -268,6 +282,7 @@ class UdopModelTest(ModelTesterMixin, PipelineTesterMixin, unittest.TestCase):
     test_resize_embeddings = True
     test_model_parallel = False
     is_encoder_decoder = True
+    test_cpu_offload = False
     # The small UDOP model needs higher percentages for CPU/MP tests
     model_split_percents = [0.8, 0.9]
 
@@ -491,10 +506,11 @@ class UdopEncoderOnlyModelTester:
         self,
         config,
         input_ids,
+        bbox,
         attention_mask,
     ):
         model = UdopEncoderModel(config=config).to(torch_device).half().eval()
-        output = model(input_ids, attention_mask=attention_mask)["last_hidden_state"]
+        output = model(input_ids, bbox=bbox, attention_mask=attention_mask)["last_hidden_state"]
         self.parent.assertFalse(torch.isnan(output).any().item())
 
 
@@ -504,7 +520,7 @@ class UdopEncoderOnlyModelTest(ModelTesterMixin, unittest.TestCase):
     test_torchscript = False
     test_head_masking = False
     test_resize_embeddings = False
-    test_model_parallel = True
+    test_model_parallel = False
     all_parallelizable_model_classes = (UdopEncoderModel,) if is_torch_available() else ()
 
     def setUp(self):
@@ -518,11 +534,6 @@ class UdopEncoderOnlyModelTest(ModelTesterMixin, unittest.TestCase):
         config_and_inputs = self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_model(*config_and_inputs)
 
-    @unittest.skipIf(torch_device == "cpu", "Cant do half precision")
-    def test_model_fp16_forward(self):
-        config_and_inputs = self.model_tester.prepare_config_and_inputs()
-        self.model_tester.create_and_check_model_fp16_forward(*config_and_inputs)
-
     @unittest.skip(
         "Not currently compatible. Fails with - NotImplementedError: Cannot copy out of meta tensor; no data!"
     )
@@ -558,7 +569,7 @@ class UdopModelIntegrationTests(unittest.TestCase):
         model = self.model
 
         prompt = "Question answering. In which year is the report made?"
-        encoding = processor(images=self.image, text=prompt, return_tensors="pt")
+        encoding = processor(images=self.image, text=prompt, return_tensors="pt").to(torch_device)
 
         predicted_ids = model.generate(**encoding)