fix tests

mosaicml · Oct 22, 2024 · 3a7c025 · 3a7c025
1 parent bd749ea
commit 3a7c025
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/tests/trainer/test_checkpoint.py b/tests/trainer/test_checkpoint.py
@@ -882,7 +882,8 @@ def _get_tmp_dir(self):
 
                     if delete_local:
                         # delete files locally, forcing trainer to look in object store
-                        shutil.rmtree('first')
+                        assert trainer_1._checkpoint_saver is not None
+                        shutil.rmtree(trainer_1._checkpoint_saver.folder)
 
                     trainer_2 = self.get_trainer(
                         latest_filename=latest_filename,