From 1e3daea3bb80b9205be8d79ae1f21ca04ee1cffe Mon Sep 17 00:00:00 2001
From: Anchen <li.anchen.au@gmail.com>
Date: Mon, 4 Mar 2024 01:07:45 +1100
Subject: [PATCH] chore(mlx-lm): add missing model_type for starcoder2 (#522)

---
 llms/mlx_lm/models/starcoder2.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/llms/mlx_lm/models/starcoder2.py b/llms/mlx_lm/models/starcoder2.py
index aeebfc9..27a53af 100644
--- a/llms/mlx_lm/models/starcoder2.py
+++ b/llms/mlx_lm/models/starcoder2.py
@@ -167,6 +167,7 @@ class Starcoder2Model(nn.Module):
 class Model(nn.Module):
     def __init__(self, args: ModelArgs):
         super().__init__()
+        self.model_type = args.model_type
         self.model = Starcoder2Model(args)
         # This is for 15B starcoder2 since it doesn't tie word embeddings
         if not args.tie_word_embeddings: