From 5278143c1b8aeff1bf19826c9b9219e5ce790d01 Mon Sep 17 00:00:00 2001
From: Angelos Katharopoulos <a_katharopoulos@apple.com>
Date: Tue, 14 Jan 2025 17:14:44 -0800
Subject: [PATCH] Fix test

---
 llms/tests/test_finetune.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/llms/tests/test_finetune.py b/llms/tests/test_finetune.py
index 6ba81628..a6d53747 100644
--- a/llms/tests/test_finetune.py
+++ b/llms/tests/test_finetune.py
@@ -21,7 +21,7 @@ from mlx_lm.tuner.utils import build_schedule
 @contextmanager
 def swapped_with_identity(obj, func):
     old_func = getattr(obj, func)
-    setattr(obj, func, lambda x: x)
+    setattr(obj, func, lambda x, **kwargs: x)
     yield
     setattr(obj, func, old_func)