Tests: improve memory usage (#1147)

bed0860b · Matthew Douglas · GitHub · 494de206 · bed0860b
Unverified Commit bed0860b authored Apr 02, 2024 by Matthew Douglas Committed by GitHub Apr 02, 2024
Hide whitespace changes
Inline Side-by-side

Showing with 9 additions and 0 deletions

tests/conftest.py tests/conftest.py +9 -0

No files found.
--- a/tests/conftest.py
+++ b/tests/conftest.py
+import gc
+
 import pytest
 import torch

@@ -20,6 +22,13 @@ def pytest_runtest_call(item):
        raise


+@pytest.hookimpl(trylast=True)
+def pytest_runtest_teardown(item, nextitem):
+    gc.collect()
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+
+
 @pytest.fixture(scope="session")
 def requires_cuda() -> bool:
    cuda_available = torch.cuda.is_available()