Add missing yield in fixture

Signed-off-by: Fabrice Normandin <[email protected]>
mila-iqia · Nov 21, 2024 · e2a18e0 · e2a18e0
1 parent 9ed0634
commit e2a18e0
Showing 1 changed file with 6 additions and 0 deletions.
diff --git 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: