From 494cdf8e96b353bbc1340765f5d8fb4660c1c9e3 Mon Sep 17 00:00:00 2001
From: Anchen <li.anchen.au@gmail.com>
Date: Sun, 24 Mar 2024 01:22:11 +1100
Subject: [PATCH] chore: fix loar for moe model (#608)

---
 llms/mlx_lm/tuner/trainer.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/llms/mlx_lm/tuner/trainer.py b/llms/mlx_lm/tuner/trainer.py
index 9bd572e3..f0d8e0a4 100644
--- a/llms/mlx_lm/tuner/trainer.py
+++ b/llms/mlx_lm/tuner/trainer.py
@@ -189,7 +189,6 @@ def train(
 
     state = [model.state, optimizer.state]
 
-    @partial(mx.compile, inputs=state, outputs=state)
     def step(batch):
         # Forward and backward pass
         (lvalue, toks), grad = loss_value_and_grad(model, *batch)