SUMIN/HuggingFace_transformer
1
0
Fork 0
Code Issues Pull Requests Actions 7 Packages Projects Releases Wiki Activity

Fix GitModelIntegrationTest.test_batched_generation device issue (#21362)

Browse Source 
fix

Co-authored-by: ydshieh <ydshieh@users.noreply.github.com>
This commit is contained in:
Yih-Dar
2023-01-30 10:37:56 +01:00
committed by GitHub
parent 73a2ff6974
commit a582cfce3c
1 changed files with 2 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
tests/models/git/test_modeling_git.py
View File

@@ -508,9 +508,8 @@ class GitModelIntegrationTest(unittest.TestCase):
# we have to prepare `input_ids` with the same batch size as `pixel_values` # we have to prepare `input_ids` with the same batch size as `pixel_values`
start_token_id = model.config.bos_token_id start_token_id = model.config.bos_token_id
generated_ids = model.generate( input_ids = torch.tensor([[start_token_id], [start_token_id]], device=torch_device)
pixel_values=pixel_values, input_ids=torch.tensor([[start_token_id], [start_token_id]]), max_length=50 generated_ids = model.generate(pixel_values=pixel_values, input_ids=input_ids, max_length=50)
)
generated_captions = processor.batch_decode(generated_ids, skip_special_tokens=True) generated_captions = processor.batch_decode(generated_ids, skip_special_tokens=True)
self.assertEquals(generated_captions, ["two cats sleeping on a pink blanket next to remotes."] * 2) self.assertEquals(generated_captions, ["two cats sleeping on a pink blanket next to remotes."] * 2)