[SAM] Fix sam slow test (#24140)

* fix sam test * update pipeline typehint
2023-06-09 16:22:09 +02:00
parent 847b47c0ee
commit 62fe753325
2 changed files with 4 additions and 2 deletions
--- a/tests/models/sam/test_modeling_sam.py
+++ b/tests/models/sam/test_modeling_sam.py
@@ -760,7 +760,9 @@ class SamModelIntegrationTest(unittest.TestCase):
        torch.testing.assert_allclose(iou_scores, EXPECTED_IOU, atol=1e-4, rtol=1e-4)

    def test_dummy_pipeline_generation(self):
-        generator = pipeline("mask-generation", model="facebook/sam-vit-base", device=torch_device)
+        generator = pipeline(
+            "mask-generation", model="facebook/sam-vit-base", device=0 if torch.cuda.is_available() else -1
+        )
        raw_image = prepare_image()

        _ = generator(raw_image, points_per_batch=64)