From efa5a6362458846952ff33154f4ce13027ba436b Mon Sep 17 00:00:00 2001 From: "Adam J. Stewart" Date: Sat, 31 Aug 2024 14:37:53 +0200 Subject: [PATCH] Automatically delete downloaded models after successful tests to avoid OOM --- tests/conftest.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/tests/conftest.py b/tests/conftest.py index 7d622c3977c..4d5a43c810d 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -1,6 +1,7 @@ # Copyright (c) Microsoft Corporation. All rights reserved. # Licensed under the MIT License. +from pathlib import Path from typing import Any import pytest @@ -16,3 +17,8 @@ def load(*args: Any, progress: bool = False, **kwargs: Any) -> Any: @pytest.fixture def load_state_dict_from_url(monkeypatch: MonkeyPatch) -> None: monkeypatch.setattr(torchvision.models._api, 'load_state_dict_from_url', load) + + +@pytest.fixture(autouse=True) +def torch_hub(tmp_path: Path) -> None: + torch.hub.set_dir(tmp_path)