fix vit hybrid test (#25543)

fix test

fix vit hybrid test (#25543)
fix test
0ed23e4d · Marc Sun · GitHub · 3f9cb335 · 0ed23e4d
Unverified Commit 0ed23e4d authored Aug 16, 2023 by Marc Sun Committed by GitHub Aug 16, 2023
Hide whitespace changes
Inline Side-by-side

Showing with 1 addition and 1 deletion

tests/models/vit_hybrid/test_modeling_vit_hybrid.py tests/models/vit_hybrid/test_modeling_vit_hybrid.py +1 -1

No files found.
--- 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