Skip to content

GitLab

Unverified Commit 62fe7533 authored by Younes Belkada's avatar Younes Belkada Committed by GitHub
Browse files

[`SAM`] Fix sam slow test (#24140) 

* fix sam test

* update pipeline typehint
parent 847b47c0
Hide whitespace changes
Inline Side-by-side
Showing with 4 additions and 2 deletions
src/transformers/pipelines/base.py
View file @ 62fe7533
...@@ -760,7 +760,7 @@ class Pipeline(_ScikitCompat): ...@@ -760,7 +760,7 @@ class Pipeline(_ScikitCompat):
framework: Optional[str] = None, framework: Optional[str] = None,
task: str = "", task: str = "",
args_parser: ArgumentHandler = None, args_parser: ArgumentHandler = None,
device: Union[int, str, "torch.device"] = None, device: Union[int, "torch.device"] = None,
torch_dtype: Optional[Union[str, "torch.dtype"]] = None, torch_dtype: Optional[Union[str, "torch.dtype"]] = None,
binary_output: bool = False, binary_output: bool = False,
**kwargs, **kwargs,
......
tests/models/sam/test_modeling_sam.py
View file @ 62fe7533
...@@ -760,7 +760,9 @@ class SamModelIntegrationTest(unittest.TestCase): ...@@ -760,7 +760,9 @@ class SamModelIntegrationTest(unittest.TestCase):
torch.testing.assert_allclose(iou_scores, EXPECTED_IOU, atol=1e-4, rtol=1e-4) torch.testing.assert_allclose(iou_scores, EXPECTED_IOU, atol=1e-4, rtol=1e-4)
def test_dummy_pipeline_generation(self): 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() raw_image = prepare_image()
_ = generator(raw_image, points_per_batch=64) _ = generator(raw_image, points_per_batch=64)
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment