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

[tests] make 2 tests device-agnostic (#30008)

Browse Source 
add torch device
This commit is contained in:
Fanli Lin
2024-04-10 20:46:39 +08:00
committed by GitHub
parent bb76f81e40
commit 185463784e
2 changed files with 3 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
tests/models/blip_2/test_modeling_blip_2.py
View File

@@ -992,7 +992,7 @@ class Blip2ModelIntegrationTest(unittest.TestCase):
# prepare image
image = prepare_img()
inputs = processor(images=image, return_tensors="pt").to(0, dtype=torch.float16)
inputs = processor(images=image, return_tensors="pt").to(f"{torch_device}:0", dtype=torch.float16)
predictions = model.generate(**inputs)
generated_text = processor.batch_decode(predictions, skip_special_tokens=True)[0].strip()
@@ -1003,7 +1003,7 @@ class Blip2ModelIntegrationTest(unittest.TestCase):
# image and context
prompt = "Question: which city is this? Answer:"
inputs = processor(images=image, text=prompt, return_tensors="pt").to(0, dtype=torch.float16)
inputs = processor(images=image, text=prompt, return_tensors="pt").to(f"{torch_device}:0", dtype=torch.float16)
predictions = model.generate(**inputs)
generated_text = processor.batch_decode(predictions, skip_special_tokens=True)[0].strip()