New checkpointing format in flax

allenai · Apr 5, 2024 · fbcc50c · fbcc50c
1 parent 29776d0
commit fbcc50c
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/tests/integrations/flax/train_test.py b/tests/integrations/flax/train_test.py
@@ -20,5 +20,5 @@ def test_trainer(self):
             ],
         )
         assert (
-            result_dir / "train" / "work" / "checkpoint_state_latest" / "checkpoint_0"
+            result_dir / "train" / "work" / "checkpoint_state_latest" / "checkpoint_0" / "checkpoint"
         ).is_file()