Skip to content

Commit

Permalink
Automatically delete downloaded models after successful tests to avoi…
Browse files Browse the repository at this point in the history
…d OOM
  • Loading branch information
adamjstewart committed Aug 31, 2024
1 parent a3c2777 commit efa5a63
Showing 1 changed file with 6 additions and 0 deletions.
6 changes: 6 additions & 0 deletions tests/conftest.py
Original file line number Diff line number Diff line change
@@ -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
Expand All @@ -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)

0 comments on commit efa5a63

Please sign in to comment.