From 7e1039453cc1a23616c11885260f1b6a7f450788 Mon Sep 17 00:00:00 2001
From: Zach Mueller <muellerzr@gmail.com>
Date: Wed, 24 Apr 2024 13:32:39 -0400
Subject: [PATCH] Fix failing tests

---
 tests/trainer/test_trainer_callback.py | 18 +++++++++++-------
 1 file changed, 11 insertions(+), 7 deletions(-)

diff --git a/tests/trainer/test_trainer_callback.py b/tests/trainer/test_trainer_callback.py
index e08ec5065fd891..8c0c9367d8d779 100644
--- a/tests/trainer/test_trainer_callback.py
+++ b/tests/trainer/test_trainer_callback.py
@@ -269,7 +269,7 @@ def test_stateful_callbacks(self):
             callbacks=[cb],
             load_best_model_at_end=True,
             save_strategy="steps",
-            evaluation_strategy="steps",
+            eval_strategy="steps",
             save_steps=2,
             eval_steps=2,
             max_steps=2,
@@ -281,7 +281,7 @@ def test_stateful_callbacks(self):
             callbacks=[EarlyStoppingCallback()],
             load_best_model_at_end=True,
             save_strategy="steps",
-            evaluation_strategy="steps",
+            eval_strategy="steps",
             save_steps=2,
             eval_steps=2,
             max_steps=2,
@@ -307,7 +307,7 @@ def test_stateful_mixed_callbacks(self):
             callbacks=cbs,
             load_best_model_at_end=True,
             save_strategy="steps",
-            evaluation_strategy="steps",
+            eval_strategy="steps",
             save_steps=2,
             eval_steps=2,
             max_steps=2,
@@ -319,7 +319,7 @@ def test_stateful_mixed_callbacks(self):
             callbacks=[EarlyStoppingCallback(), MyTestTrainerCallback()],
             load_best_model_at_end=True,
             save_strategy="steps",
-            evaluation_strategy="steps",
+            eval_strategy="steps",
             save_steps=2,
             eval_steps=2,
             max_steps=2,
@@ -346,7 +346,7 @@ def test_stateful_duplicate_callbacks(self):
             callbacks=cbs,
             load_best_model_at_end=True,
             save_strategy="steps",
-            evaluation_strategy="steps",
+            eval_strategy="steps",
             save_steps=2,
             eval_steps=2,
             max_steps=2,
@@ -358,7 +358,7 @@ def test_stateful_duplicate_callbacks(self):
             callbacks=[MyTestExportableCallback(), MyTestExportableCallback()],
             load_best_model_at_end=True,
             save_strategy="steps",
-            evaluation_strategy="steps",
+            eval_strategy="steps",
             save_steps=2,
             eval_steps=2,
             max_steps=2,
@@ -382,7 +382,7 @@ def test_missing_stateful_callback(self):
             callbacks=[cb],
             load_best_model_at_end=True,
             save_strategy="steps",
-            evaluation_strategy="steps",
+            eval_strategy="steps",
             save_steps=2,
             eval_steps=2,
             max_steps=2,
@@ -391,6 +391,10 @@ def test_missing_stateful_callback(self):
 
         # Create a new trainer with defaults
         trainer = self.get_trainer(
+            save_strategy="steps",
+            eval_strategy="steps",
+            save_steps=2,
+            eval_steps=2,
             max_steps=2,
             restore_callback_states_from_checkpoint=True,
         )