diff --git a/llms/mlx_lm/generate.py b/llms/mlx_lm/generate.py
index cf159c21..58ecde05 100644
--- a/llms/mlx_lm/generate.py
+++ b/llms/mlx_lm/generate.py
@@ -115,7 +115,14 @@ def main(args):
     formatter = colorprint_by_t0 if args.colorize else None
 
     generate(
-        model, tokenizer, prompt, args.temp, args.max_tokens, True, formatter=formatter, top_p=args.top_p
+        model,
+        tokenizer,
+        prompt,
+        args.temp,
+        args.max_tokens,
+        True,
+        formatter=formatter,
+        top_p=args.top_p,
     )
 
 
diff --git a/llms/mlx_lm/merge.py b/llms/mlx_lm/merge.py
index efdc726d..d2f54316 100644
--- a/llms/mlx_lm/merge.py
+++ b/llms/mlx_lm/merge.py
@@ -3,6 +3,7 @@
 import argparse
 import glob
 import json
+import shutil
 from pathlib import Path
 
 import mlx.core as mx