Skip to content

Commit d39cfec

Browse files
authored
fix: add sanitize method to Granite model for tied embeddings (ml-explore#1298)
1 parent 04a1910 commit d39cfec

1 file changed

Lines changed: 5 additions & 0 deletions

File tree

mlx_lm/models/granite.py

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -191,3 +191,8 @@ def __call__(
191191
@property
192192
def layers(self):
193193
return self.model.layers
194+
195+
def sanitize(self, weights):
196+
if self.args.tie_word_embeddings:
197+
weights.pop("lm_head.weight", None)
198+
return weights

0 commit comments

Comments
 (0)