From e2a18e0c1ddd074b9b68a47d5c8da8727f530034 Mon Sep 17 00:00:00 2001 From: Fabrice Normandin Date: Thu, 21 Nov 2024 19:14:09 +0000 Subject: [PATCH] Add missing `yield` in fixture Signed-off-by: Fabrice Normandin --- project/conftest.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/project/conftest.py b/project/conftest.py index 6ed38396..a30e800a 100644 --- a/project/conftest.py +++ b/project/conftest.py @@ -144,6 +144,12 @@ def prevent_jax_from_reserving_all_the_vram(): # note; not using monkeypatch because we want this to be session-scoped. val_before = os.environ.get("XLA_PYTHON_CLIENT_PREALLOCATE") os.environ["XLA_PYTHON_CLIENT_PREALLOCATE"] = "false" + + # allocator_before = os.environ.get("XLA_PYTHON_CLIENT_ALLOCATOR") + # os.environ["XLA_PYTHON_CLIENT_ALLOCATOR"] = "platform" + + yield + if val_before is None: os.environ.pop("XLA_PYTHON_CLIENT_PREALLOCATE") else: