[LayoutLMV2 Tests] Make sure input is on GPU (#15314)

* [LayoutLMV2 Tests] Make sure input is on GPU * correct empty line

[LayoutLMV2 Tests] Make sure input is on GPU (#15314)
* [LayoutLMV2 Tests] Make sure input is on GPU * correct empty line
dcaa5100 · Patrick von Platen · GitHub · c15bb3fe · dcaa5100
Unverified Commit dcaa5100 authored Jan 24, 2022 by Patrick von Platen Committed by GitHub Jan 24, 2022
Hide whitespace changes
Inline Side-by-side

Showing with 2 additions and 1 deletion

tests/test_modeling_layoutlmv2.py tests/test_modeling_layoutlmv2.py +2 -1

No files found.
--- a/tests/test_modeling_layoutlmv2.py
+++ b/tests/test_modeling_layoutlmv2.py
@@ -122,7 +122,8 @@ class LayoutLMv2ModelTester:
                    bbox[i, j, 0] = t

        image = ImageList(
-            torch.zeros(self.batch_size, self.num_channels, self.image_size, self.image_size), self.image_size
+            torch.zeros(self.batch_size, self.num_channels, self.image_size, self.image_size, device=torch_device),
+            self.image_size,
        )

        input_mask = None