[Tests] Fix test (huggingface#15324)

NielsRogge · web-flow · commit 637e81752aad · 2022-01-25T15:48:25.000+01:00
* Fix Swin device

* Remove print statement
diff --git a/src/transformers/models/swin/modeling_swin.py b/src/transformers/models/swin/modeling_swin.py
@@ -437,6 +437,9 @@ def forward(self, hidden_states, head_mask=None, output_attentions=False):
         hidden_states_windows = window_partition(shifted_hidden_states, self.window_size)
         hidden_states_windows = hidden_states_windows.view(-1, self.window_size * self.window_size, channels)
 
+        if self.attn_mask is not None:
+            self.attn_mask = self.attn_mask.to(hidden_states_windows.device)
+
         self_attention_outputs = self.attention(
             hidden_states_windows,
             self.attn_mask,
diff --git a/tests/test_modeling_vilt.py b/tests/test_modeling_vilt.py
@@ -595,8 +595,8 @@ def test_inference_natural_language_visual_reasoning(self):
 
         # forward pass
         outputs = model(
-            input_ids=encoding_1.input_ids,
-            pixel_values=pixel_values,
+            input_ids=encoding_1.input_ids.to(torch_device),
+            pixel_values=pixel_values.to(torch_device),
         )
 
         # verify the logits
diff --git a/tests/test_modeling_vit_mae.py b/tests/test_modeling_vit_mae.py
@@ -327,9 +327,6 @@ def test_save_load(self):
         config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()
 
         for model_class in self.all_model_classes:
-
-            print("Model class:", model_class)
-
             model = model_class(config)
             model.to(torch_device)
             model.eval()

Original file line number	Diff line number	Diff line change
`@@ -595,8 +595,8 @@ def test_inference_natural_language_visual_reasoning(self):`
`595`	`595`
`596`	`596`	`# forward pass`
`597`	`597`	`outputs = model(`
`598`		`- input_ids=encoding_1.input_ids,`
`599`		`- pixel_values=pixel_values,`
	`598`	`+ input_ids=encoding_1.input_ids.to(torch_device),`
	`599`	`+ pixel_values=pixel_values.to(torch_device),`
`600`	`600`	`)`
`601`	`601`
`602`	`602`	`# verify the logits`