mirror of
https://github.com/ml-explore/mlx-examples.git
synced 2025-06-24 01:17:28 +08:00
remove lm head if unused (#1324)
This commit is contained in:
parent
877d2a345b
commit
595f5da146
@ -196,9 +196,12 @@ class Model(nn.Module):
|
||||
|
||||
def sanitize(self, weights):
|
||||
# Remove unused precomputed rotary freqs
|
||||
return {
|
||||
weights = {
|
||||
k: v for k, v in weights.items() if "self_attn.rotary_emb.inv_freq" not in k
|
||||
}
|
||||
if self.args.tie_word_embeddings:
|
||||
weights.pop("lm_head.weight", None)
|
||||
return weights
|
||||
|
||||
@property
|
||||
def layers(self):
|
||||
|
Loading…
Reference in New Issue
Block a user