From dcaa5100c9f35b72bb2af5538dd5e2b54858ace2 Mon Sep 17 00:00:00 2001 From: Patrick von Platen Date: Mon, 24 Jan 2022 15:54:47 +0100 Subject: [PATCH] [LayoutLMV2 Tests] Make sure input is on GPU (#15314) * [LayoutLMV2 Tests] Make sure input is on GPU * correct empty line --- tests/test_modeling_layoutlmv2.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tests/test_modeling_layoutlmv2.py b/tests/test_modeling_layoutlmv2.py index 94e274e8a1..4e02f985f7 100644 --- 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