Fix tests (#14289)

huggingface · Nov 6, 2021 · 34307bb · 34307bb
1 parent 24b30d4
commit 34307bb
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 3 deletions.
diff --git a/tests/test_modeling_beit.py b/tests/test_modeling_beit.py
@@ -232,7 +232,9 @@ def test_training(self):
             # 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 @@ def test_training_gradient_checkpointing(self):
             # 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()

diff --git a/tests/test_modeling_segformer.py b/tests/test_modeling_segformer.py
@@ -318,7 +318,9 @@ def test_training(self):
             # 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()