Fix tests (#14289)

tests/test_modeling_beit.py

@@ -232,7 +232,9 @@ class BeitModelTest(ModelTesterMixin, unittest.TestCase):
             # this can then be incorporated into _prepare_for_class in test_modeling_common.py
             elif model_class.__name__ == "BeitForSemanticSegmentation":
                 batch_size, num_channels, height, width = inputs_dict["pixel_values"].shape
-                inputs_dict["labels"] = torch.zeros([self.model_tester.batch_size, height, width]).long()
+                inputs_dict["labels"] = torch.zeros(
+                    [self.model_tester.batch_size, height, width], device=torch_device
+                ).long()
             model = model_class(config)
             model.to(torch_device)
             model.train()
@@ -259,7 +261,9 @@ class BeitModelTest(ModelTesterMixin, unittest.TestCase):
             # this can then be incorporated into _prepare_for_class in test_modeling_common.py
             elif model_class.__name__ == "BeitForSemanticSegmentation":
                 batch_size, num_channels, height, width = inputs_dict["pixel_values"].shape
-                inputs_dict["labels"] = torch.zeros([self.model_tester.batch_size, height, width]).long()
+                inputs_dict["labels"] = torch.zeros(
+                    [self.model_tester.batch_size, height, width], device=torch_device
+                ).long()
             model = model_class(config)
             model.to(torch_device)
             model.train()

tests/test_modeling_segformer.py

@@ -318,7 +318,9 @@ class SegformerModelTest(ModelTesterMixin, unittest.TestCase):
             # this can then be incorporated into _prepare_for_class in test_modeling_common.py
             if model_class.__name__ == "SegformerForSemanticSegmentation":
                 batch_size, num_channels, height, width = inputs_dict["pixel_values"].shape
-                inputs_dict["labels"] = torch.zeros([self.model_tester.batch_size, height, width]).long()
+                inputs_dict["labels"] = torch.zeros(
+                    [self.model_tester.batch_size, height, width], device=torch_device
+                ).long()
             model = model_class(config)
             model.to(torch_device)
             model.train()