From 34940c5607646a6d373e05498bf2fc352862f859 Mon Sep 17 00:00:00 2001
From: Awni Hannun <awni@apple.com>
Date: Mon, 24 Mar 2025 16:25:15 -0700
Subject: [PATCH] fix test

---
 mlx_lm/tuner/trainer.py |  5 ++---
 tests/test_finetune.py  | 15 +++++++--------
 2 files changed, 9 insertions(+), 11 deletions(-)

diff --git a/mlx_lm/tuner/trainer.py b/mlx_lm/tuner/trainer.py
index 4a71fc1..22f0908 100644
--- a/mlx_lm/tuner/trainer.py
+++ b/mlx_lm/tuner/trainer.py
@@ -15,8 +15,8 @@ from mlx.nn.utils import average_gradients
 from mlx.utils import tree_flatten, tree_map
 from transformers import PreTrainedTokenizer
 
-from .datasets import CacheDataset
 from ..models.cache import KVCache, make_prompt_cache
+from .datasets import CacheDataset
 
 
 def reset_prompt_cache(cache):
@@ -267,6 +267,7 @@ def train(
 
     model.train()
     seq_step_size = args.seq_step_size or args.max_seq_length
+
     def seq_split_step(batch):
         losses = mx.array(0.0)
         n_tokens = mx.array(0.0)
@@ -299,7 +300,6 @@ def train(
 
     loss_value_and_grad = nn.value_and_grad(model, loss)
 
->>>>>>> 568a8d6 (use gradient accumulation)
     losses = 0
     n_tokens = 0
     steps = 0
@@ -321,7 +321,6 @@ def train(
         # is always measured before any training.
         if it == 1 or it % args.steps_per_eval == 0 or it == args.iters:
             tic = time.perf_counter()
-            val_loss = 0.0
             val_loss = evaluate(
                 model=model,
                 dataset=val_dataset,
diff --git a/tests/test_finetune.py b/tests/test_finetune.py
index dd6ee8c..ca8867b 100644
--- a/tests/test_finetune.py
+++ b/tests/test_finetune.py
@@ -370,11 +370,10 @@ class TestScheduleConfig(unittest.TestCase):
         mock_iterate_batches = MagicMock()
 
         mock_iterate_batches.return_value = [
-            (MagicMock(), MagicMock()),
-            (MagicMock(), MagicMock()),
-            (MagicMock(), MagicMock()),
-            (MagicMock(), MagicMock()),
-            (MagicMock(), MagicMock()),
+            (mx.ones((2, 8), mx.int32), mx.ones((2, 2), mx.int32)),
+            (mx.ones((2, 8), mx.int32), mx.ones((2, 2), mx.int32)),
+            (mx.ones((2, 8), mx.int32), mx.ones((2, 2), mx.int32)),
+            (mx.ones((2, 8), mx.int32), mx.ones((2, 2), mx.int32)),
         ]
 
         mock_default_loss.side_effect = [
@@ -412,9 +411,9 @@ class TestScheduleConfig(unittest.TestCase):
         mock_iterate_batches = MagicMock()
 
         mock_iterate_batches.return_value = [
-            (MagicMock(), MagicMock()),
-            (MagicMock(), MagicMock()),
-            (MagicMock(), MagicMock()),
+            (mx.ones((2, 8), mx.int32), mx.ones((2, 2), mx.int32)),
+            (mx.ones((2, 8), mx.int32), mx.ones((2, 2), mx.int32)),
+            (mx.ones((2, 8), mx.int32), mx.ones((2, 2), mx.int32)),
         ]
 
         mock_default_loss.side_effect = [