reset accelerate env variables after each test (#24107)

0f236050 · Sourab Mangrulkar · GitHub · 5fa0a1b2 · 0f236050
Unverified Commit 0f236050 authored Jun 08, 2023 by Sourab Mangrulkar Committed by GitHub Jun 08, 2023
Hide whitespace changes
Inline Side-by-side

Showing with 5 additions and 0 deletions

src/transformers/testing_utils.py src/transformers/testing_utils.py +5 -0

No files found.
--- a/src/transformers/testing_utils.py
+++ b/src/transformers/testing_utils.py
@@ -1339,6 +1339,11 @@ class TestCasePlus(unittest.TestCase):
            AcceleratorState._reset_state()
            PartialState._reset_state()

+            # delete all the env variables having `ACCELERATE` in them
+            for k in list(os.environ.keys()):
+                if "ACCELERATE" in k:
+                    del os.environ[k]
+

 def mockenv(**kwargs):
    """