From 1bb6c9496963af1db937d3d9a1edcbc959e4b18b Mon Sep 17 00:00:00 2001
From: Oliver Koenig <okoenig@nvidia.com>
Date: Wed, 6 Nov 2024 02:35:10 -0800
Subject: [PATCH] ADLR/megatron-lm!2304 - tests: Set flash/fused attn

---
 tests/unit_tests/conftest.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tests/unit_tests/conftest.py b/tests/unit_tests/conftest.py
index 01b5bcb30b..8fb1c3f99a 100644
--- a/tests/unit_tests/conftest.py
+++ b/tests/unit_tests/conftest.py
@@ -13,7 +13,7 @@
 from tests.unit_tests.test_utilities import Utils
 
 
-@pytest.fixture(scope="session", autouse=True)
+@pytest.fixture(scope="function", autouse=True)
 def set_env():
     if is_te_min_version("1.3"):
         os.environ['NVTE_FLASH_ATTN'] = '0'