[Pix2Struct] Fix slow test (#22448)

fix slow test
2023-03-29 17:40:45 +02:00
parent 55dae94c0c
commit b844f8a9ab
1 changed files with 1 additions and 1 deletions
--- a/tests/models/pix2struct/test_modeling_pix2struct.py
+++ b/tests/models/pix2struct/test_modeling_pix2struct.py
@@ -406,7 +406,7 @@ class Pix2StructTextImageModelTest(ModelTesterMixin, unittest.TestCase):
    def test_model(self):
        config, input_dict = self.model_tester.prepare_config_and_inputs_for_common()
        for model_class in self.all_model_classes:
-            model = model_class(config)
+            model = model_class(config).to(torch_device)

            output = model(**input_dict)
            self.assertEqual(