riaz.somc/transformers
1
0
Fork 0
mirror of https://github.com/huggingface/transformers.git synced 2025-07-31 02:02:21 +06:00
Code Issues Actions 23 Packages Projects Releases Wiki Activity

Fix SAM OOM issue on CI (#24125)

Browse Source 
fix

Co-authored-by: ydshieh <ydshieh@users.noreply.github.com>
This commit is contained in:
Yih-Dar 2023-06-09 15:07:08 +02:00 committed by GitHub
parent 707023d155
commit b8fe259f16
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 13 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

7
tests/models/sam/test_modeling_sam.py
View File

@ -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")

6
tests/models/sam/test_modeling_tf_sam.py
View File

@ -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")