call GC between each e2e test to avoid OOM

finegrain-ai · Jan 29, 2024 · 2aeb968 · 2aeb968
1 parent deb5e93
commit 2aeb968
Showing 1 changed file with 8 additions and 0 deletions.
diff --git a/tests/e2e/test_diffusion.py b/tests/e2e/test_diffusion.py
@@ -1,3 +1,4 @@
+import gc
 from pathlib import Path
 from typing import Iterator
 from warnings import warn
@@ -30,6 +31,13 @@
 from tests.utils import ensure_similar_images
 
 
+@pytest.fixture(autouse=True)
+def ensure_gc():
+    # Avoid GPU OOMs
+    # See https://github.com/pytest-dev/pytest/discussions/8153#discussioncomment-214812
+    gc.collect()
+
+
 @pytest.fixture(scope="module")
 def ref_path(test_e2e_path: Path) -> Path:
     return test_e2e_path / "test_diffusion_ref"