Unverified Commit bed0860b authored by Matthew Douglas's avatar Matthew Douglas Committed by GitHub
Browse files

Tests: improve memory usage (#1147)

parent 494de206
import gc
import pytest import pytest
import torch import torch
...@@ -20,6 +22,13 @@ def pytest_runtest_call(item): ...@@ -20,6 +22,13 @@ def pytest_runtest_call(item):
raise 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") @pytest.fixture(scope="session")
def requires_cuda() -> bool: def requires_cuda() -> bool:
cuda_available = torch.cuda.is_available() cuda_available = torch.cuda.is_available()
......
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