Fix device issue in `SwiftFormerModelIntegrationTest::test_inference_…

…image_classification_head` (huggingface#23435) fix Co-authored-by: ydshieh <[email protected]>
dkqkxx · May 17, 2023 · a8732e0 · a8732e0
1 parent 0f2c738
commit a8732e0
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/tests/models/swiftformer/test_modeling_swiftformer.py b/tests/models/swiftformer/test_modeling_swiftformer.py
@@ -303,5 +303,5 @@ def test_inference_image_classification_head(self):
         expected_shape = torch.Size((1, 1000))
         self.assertEqual(outputs.logits.shape, expected_shape)
 
-        expected_slice = torch.tensor([[-2.1703e00, 2.1107e00, -2.0811e00]])
+        expected_slice = torch.tensor([[-2.1703e00, 2.1107e00, -2.0811e00]]).to(torch_device)
         self.assertTrue(torch.allclose(outputs.logits[0, :3], expected_slice, atol=1e-4))