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

patrickvonplaten · web-flow · commit dcaa5100c9f3 · 2022-01-24T15:54:47.000+01:00
* [LayoutLMV2 Tests] Make sure input is on GPU

* correct empty line
diff --git a/tests/test_modeling_layoutlmv2.py b/tests/test_modeling_layoutlmv2.py
@@ -122,7 +122,8 @@ def prepare_config_and_inputs(self):
                     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

Original file line number	Diff line number	Diff line change
`@@ -122,7 +122,8 @@ def prepare_config_and_inputs(self):`
`122`	`122`	`bbox[i, j, 0] = t`
`123`	`123`
`124`	`124`	`image = ImageList(`
`125`		`- torch.zeros(self.batch_size, self.num_channels, self.image_size, self.image_size), self.image_size`
	`125`	`+ torch.zeros(self.batch_size, self.num_channels, self.image_size, self.image_size, device=torch_device),`
	`126`	`+ self.image_size,`
`126`	`127`	`)`
`127`	`128`
`128`	`129`	`input_mask = None`