From b8fe259f163c48a18c9b27428b72b2ac104de346 Mon Sep 17 00:00:00 2001 From: Yih-Dar <2521628+ydshieh@users.noreply.github.com> Date: Fri, 9 Jun 2023 15:07:08 +0200 Subject: [PATCH] Fix SAM OOM issue on CI (#24125) fix Co-authored-by: ydshieh --- tests/models/sam/test_modeling_sam.py | 7 +++++++ tests/models/sam/test_modeling_tf_sam.py | 6 ++++++ 2 files changed, 13 insertions(+) diff --git a/tests/models/sam/test_modeling_sam.py b/tests/models/sam/test_modeling_sam.py index a701514522..c8bc30a298 100644 --- a/tests/models/sam/test_modeling_sam.py +++ b/tests/models/sam/test_modeling_sam.py @@ -15,6 +15,7 @@ """ Testing suite for the PyTorch SAM model. """ +import gc import inspect import unittest @@ -461,6 +462,12 @@ def prepare_dog_img(): @slow class SamModelIntegrationTest(unittest.TestCase): + def tearDown(self): + super().tearDown() + # clean-up as much as possible GPU memory occupied by PyTorch + gc.collect() + torch.cuda.empty_cache() + def test_inference_mask_generation_no_point(self): model = SamModel.from_pretrained("facebook/sam-vit-base") processor = SamProcessor.from_pretrained("facebook/sam-vit-base") diff --git a/tests/models/sam/test_modeling_tf_sam.py b/tests/models/sam/test_modeling_tf_sam.py index 4e918a1cd1..33cb38c6fb 100644 --- a/tests/models/sam/test_modeling_tf_sam.py +++ b/tests/models/sam/test_modeling_tf_sam.py @@ -17,6 +17,7 @@ from __future__ import annotations +import gc import inspect import unittest @@ -429,6 +430,11 @@ def prepare_dog_img(): @slow class SamModelIntegrationTest(unittest.TestCase): + def tearDown(self): + super().tearDown() + # clean-up as much as possible GPU memory occupied by PyTorch + gc.collect() + def test_inference_mask_generation_no_point(self): model = TFSamModel.from_pretrained("facebook/sam-vit-base") processor = SamProcessor.from_pretrained("facebook/sam-vit-base")