diff --git a/tests/test_accelerator.py b/tests/test_accelerator.py index 2caf8b39fbe..a5ff3d0d50c 100644 --- 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 )