Improve test reliability for Accelerator.free_memory() (huggingface#2935

)
byi8220 · Jul 16, 2024 · d16d737 · d16d737
1 parent 7a5c231
commit d16d737
Showing 1 changed file with 5 additions and 0 deletions.
diff --git a/tests/test_accelerator.py b/tests/test_accelerator.py
@@ -15,6 +15,7 @@
 import os
 import pickle
 import tempfile
+import time
 from unittest.mock import patch
 
 import psutil
@@ -220,6 +221,10 @@ def test_free_memory_dereferences_prepared_components(self):
         model, optimizer, scheduler, train_dl, valid_dl = accelerator.prepare(
             model, optimizer, scheduler, train_dl, valid_dl
         )
+
+        # Short sleep here makes this test more reliable
+        time.sleep(1e-3)
+
         model, optimizer, scheduler, train_dl, valid_dl = accelerator.free_memory(
             model, optimizer, scheduler, train_dl, valid_dl
         )