From 424cb854e930ef0bced08e260a5eeea28bd566bc Mon Sep 17 00:00:00 2001 From: Goekdeniz-Guelmez Date: Sun, 19 Jan 2025 02:03:50 +0100 Subject: [PATCH] nits --- llms/mlx_lm/tuner/orpo_trainer.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/llms/mlx_lm/tuner/orpo_trainer.py b/llms/mlx_lm/tuner/orpo_trainer.py index 00066df9..4eaace31 100644 --- a/llms/mlx_lm/tuner/orpo_trainer.py +++ b/llms/mlx_lm/tuner/orpo_trainer.py @@ -7,7 +7,7 @@ import mlx.core as mx import numpy as np from mlx.utils import tree_flatten from mlx.nn.utils import average_gradients -from .trainer import TrainingArgs, grad_checkpoint +from .trainer import TrainingArgs, grad_checkpoint, TrainingCallback @dataclass