fix vit hybrid test (#25543)

fix test
2023-08-16 11:02:57 -04:00 · 2023-08-16 11:02:57 -04:00 · 0ed23e4db2
parent 3f9cb33504
commit 0ed23e4db2
1 changed files with 1 additions and 1 deletions
--- a/tests/models/vit_hybrid/test_modeling_vit_hybrid.py
+++ b/tests/models/vit_hybrid/test_modeling_vit_hybrid.py
@ -281,7 +281,7 @@ class ViTModelIntegrationTest(unittest.TestCase):

        image = prepare_img()

-        inputs = image_processor(images=image, return_tensors="pt")
+        inputs = image_processor(images=image, return_tensors="pt").to(torch_device)
        outputs = model(**inputs)
        logits = outputs.logits
        # model predicts one of the 1000 ImageNet classes