From 5a551df92bcb9c8cdcd706a4458764eabc1e0fe8 Mon Sep 17 00:00:00 2001 From: Yih-Dar <2521628+ydshieh@users.noreply.github.com> Date: Mon, 4 Dec 2023 09:40:42 +0100 Subject: [PATCH] Fix `TvpModelIntegrationTests` (#27792) * fix * fix * fix --------- Co-authored-by: ydshieh --- tests/models/tvp/test_modeling_tvp.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/tests/models/tvp/test_modeling_tvp.py b/tests/models/tvp/test_modeling_tvp.py index ebdd4fb0b7..14ec02ed6f 100644 --- a/tests/models/tvp/test_modeling_tvp.py +++ b/tests/models/tvp/test_modeling_tvp.py @@ -230,10 +230,11 @@ class TvpModelIntegrationTests(unittest.TestCase): image_processor = self.default_image_processor image = prepare_img() - encoding = image_processor(images=image, return_tensors="pt").to(torch_device) + encoding = image_processor(images=image, return_tensors="pt") input_ids = torch.tensor([[1, 2]]) attention_mask = torch.tensor([[1, 1]]) encoding.update({"input_ids": input_ids, "attention_mask": attention_mask}) + encoding.to(torch_device) with torch.no_grad(): outputs = model(**encoding) @@ -250,10 +251,11 @@ class TvpModelIntegrationTests(unittest.TestCase): image_processor = self.default_image_processor image = prepare_img() - encoding = image_processor(images=image, return_tensors="pt").to(torch_device) + encoding = image_processor(images=image, return_tensors="pt") input_ids = torch.tensor([[1, 2]]) attention_mask = torch.tensor([[1, 1]]) encoding.update({"input_ids": input_ids, "attention_mask": attention_mask}) + encoding.to(torch_device) with torch.no_grad(): outputs = model(**encoding)