Fix TvpModelIntegrationTests (#27792)

* fix

* fix

* fix

---------

Co-authored-by: ydshieh <ydshieh@users.noreply.github.com>
This commit is contained in:
Yih-Dar
2023-12-04 09:40:42 +01:00
committed by GitHub
parent c0b9db0914
commit 5a551df92b

View File

@@ -230,10 +230,11 @@ class TvpModelIntegrationTests(unittest.TestCase):
image_processor = self.default_image_processor image_processor = self.default_image_processor
image = prepare_img() 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]]) input_ids = torch.tensor([[1, 2]])
attention_mask = torch.tensor([[1, 1]]) attention_mask = torch.tensor([[1, 1]])
encoding.update({"input_ids": input_ids, "attention_mask": attention_mask}) encoding.update({"input_ids": input_ids, "attention_mask": attention_mask})
encoding.to(torch_device)
with torch.no_grad(): with torch.no_grad():
outputs = model(**encoding) outputs = model(**encoding)
@@ -250,10 +251,11 @@ class TvpModelIntegrationTests(unittest.TestCase):
image_processor = self.default_image_processor image_processor = self.default_image_processor
image = prepare_img() 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]]) input_ids = torch.tensor([[1, 2]])
attention_mask = torch.tensor([[1, 1]]) attention_mask = torch.tensor([[1, 1]])
encoding.update({"input_ids": input_ids, "attention_mask": attention_mask}) encoding.update({"input_ids": input_ids, "attention_mask": attention_mask})
encoding.to(torch_device)
with torch.no_grad(): with torch.no_grad():
outputs = model(**encoding) outputs = model(**encoding)